Precise return type for Result::rowCount() based on detected driver

extension.neon

@@ -313,6 +313,18 @@ services:
 		class: PHPStan\Type\Doctrine\DBAL\QueryBuilder\QueryBuilderExecuteMethodExtension
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: PHPStan\Type\Doctrine\DBAL\RowCountMethodDynamicReturnTypeExtension
+		arguments:
+			class: Doctrine\DBAL\Result
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: PHPStan\Type\Doctrine\DBAL\RowCountMethodDynamicReturnTypeExtension
+		arguments:
+			class: Doctrine\DBAL\Driver\Result
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension
 
 	# Type descriptors
 	-

src/Type/Doctrine/DBAL/RowCountMethodDynamicReturnTypeExtension.php

@@ -0,0 +1,118 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\DBAL;
+
+use Doctrine\DBAL\Driver\Result as DriverResult;
+use Doctrine\ORM\EntityManagerInterface;
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Doctrine\Driver\DriverDetector;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Reflection\ReflectionProvider;
+use PHPStan\Type\Doctrine\ObjectMetadataResolver;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\Type;
+
+class RowCountMethodDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+
+	/** @var string */
+	private $class;
+
+	/** @var ObjectMetadataResolver */
+	private $objectMetadataResolver;
+
+	/** @var DriverDetector */
+	private $driverDetector;
+
+	/** @var ReflectionProvider */
+	private $reflectionProvider;
+
+	public function __construct(
+		string $class,
+		ObjectMetadataResolver $objectMetadataResolver,
+		DriverDetector $driverDetector,
+		ReflectionProvider $reflectionProvider
+	)
+	{
+		$this->class = $class;
+		$this->objectMetadataResolver = $objectMetadataResolver;
+		$this->driverDetector = $driverDetector;
+		$this->reflectionProvider = $reflectionProvider;
+	}
+
+	public function getClass(): string
+	{
+		return $this->class;
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		return $methodReflection->getName() === 'rowCount';
+	}
+
+	public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): ?Type
+	{
+		$objectManager = $this->objectMetadataResolver->getObjectManager();
+		if (!$objectManager instanceof EntityManagerInterface) {
+			return null;
+		}
+
+		$connection = $objectManager->getConnection();
+		$driver = $this->driverDetector->detect($connection);
+		if ($driver === null) {
+			return null;
+		}
+
+		$resultClass = $this->getResultClass($driver);
+		if ($resultClass === null) {
+			return null;
+		}
+
+		if (!$this->reflectionProvider->hasClass($resultClass)) {
+			return null;
+		}
+
+		$resultReflection = $this->reflectionProvider->getClass($resultClass);
+		if (!$resultReflection->hasNativeMethod('rowCount')) {
+			return null;
+		}
+
+		$rowCountMethod = $resultReflection->getNativeMethod('rowCount');
+		$variant = ParametersAcceptorSelector::selectSingle($rowCountMethod->getVariants());
+
+		return $variant->getReturnType();
+	}
+
+	/**
+	 * @param DriverDetector::* $driver
+	 * @return class-string<DriverResult>|null
+	 */
+	private function getResultClass(string $driver): ?string
+	{
+		switch ($driver) {
+			case DriverDetector::IBM_DB2:
+				return 'Doctrine\DBAL\Driver\IBMDB2\Result';
+			case DriverDetector::MYSQLI:
+				return 'Doctrine\DBAL\Driver\Mysqli\Result';
+			case DriverDetector::OCI8:
+				return 'Doctrine\DBAL\Driver\OCI8\Result';
+			case DriverDetector::PDO_MYSQL:
+			case DriverDetector::PDO_OCI:
+			case DriverDetector::PDO_PGSQL:
+			case DriverDetector::PDO_SQLITE:
+			case DriverDetector::PDO_SQLSRV:
+				return 'Doctrine\DBAL\Driver\PDO\Result';
+			case DriverDetector::PGSQL:
+				return 'Doctrine\DBAL\Driver\PgSQL\Result';
+			case DriverDetector::SQLITE3:
+				return 'Doctrine\DBAL\Driver\SQLite3\Result';
+			case DriverDetector::SQLSRV:
+				return 'Doctrine\DBAL\Driver\SQLSrv\Result';
+		}
+
+		return null;
+	}
+
+}

tests/Type/Doctrine/DBAL/MysqliResultRowCountReturnTypeTest.php

@@ -0,0 +1,35 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\DBAL;
+
+use PHPStan\Testing\TypeInferenceTestCase;
+
+class MysqliResultRowCountReturnTypeTest extends TypeInferenceTestCase
+{
+
+	/** @return iterable<mixed> */
+	public function dataFileAsserts(): iterable
+	{
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/mysqli-result-row-count.php');
+	}
+
+	/**
+	 * @dataProvider dataFileAsserts
+	 * @param mixed ...$args
+	 */
+	public function testFileAsserts(
+		string $assertType,
+		string $file,
+		...$args
+	): void
+	{
+		$this->assertFileAsserts($assertType, $file, ...$args);
+	}
+
+	/** @return string[] */
+	public static function getAdditionalConfigFiles(): array
+	{
+		return [__DIR__ . '/mysqli.neon'];
+	}
+
+}

tests/Type/Doctrine/DBAL/PDOResultRowCountReturnTypeTest.php

@@ -0,0 +1,35 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\DBAL;
+
+use PHPStan\Testing\TypeInferenceTestCase;
+
+class PDOResultRowCountReturnTypeTest extends TypeInferenceTestCase
+{
+
+	/** @return iterable<mixed> */
+	public function dataFileAsserts(): iterable
+	{
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/pdo-result-row-count.php');
+	}
+
+	/**
+	 * @dataProvider dataFileAsserts
+	 * @param mixed ...$args
+	 */
+	public function testFileAsserts(
+		string $assertType,
+		string $file,
+		...$args
+	): void
+	{
+		$this->assertFileAsserts($assertType, $file, ...$args);
+	}
+
+	/** @return string[] */
+	public static function getAdditionalConfigFiles(): array
+	{
+		return [__DIR__ . '/pdo.neon'];
+	}
+
+}

tests/Type/Doctrine/DBAL/data/mysqli-result-row-count.php

@@ -0,0 +1,15 @@
+<?php
+
+namespace MysqliResultRowCount;
+
+use Doctrine\DBAL\Result;
+use Doctrine\DBAL\Driver\Result as DriverResult;
+use function PHPStan\Testing\assertType;
+
+function (Result $r): void {
+	assertType('int|numeric-string', $r->rowCount());
+};
+
+function (DriverResult $r): void {
+	assertType('int|numeric-string', $r->rowCount());
+};

tests/Type/Doctrine/DBAL/data/pdo-result-row-count.php

@@ -0,0 +1,15 @@
+<?php
+
+namespace PDOResultRowCount;
+
+use Doctrine\DBAL\Result;
+use Doctrine\DBAL\Driver\Result as DriverResult;
+use function PHPStan\Testing\assertType;
+
+function (Result $r): void {
+	assertType('int', $r->rowCount());
+};
+
+function (DriverResult $r): void {
+	assertType('int', $r->rowCount());
+};

tests/Type/Doctrine/DBAL/mysqli.neon

@@ -0,0 +1,6 @@
+includes:
+	- ../../../../extension.neon
+
+parameters:
+	doctrine:
+		objectManagerLoader: mysqli.php

tests/Type/Doctrine/DBAL/mysqli.php

@@ -0,0 +1,25 @@
+<?php declare(strict_types = 1);
+
+use Cache\Adapter\PHPArray\ArrayCachePool;
+use Doctrine\Common\Annotations\AnnotationReader;
+use Doctrine\DBAL\DriverManager;
+use Doctrine\ORM\Configuration;
+use Doctrine\ORM\EntityManager;
+use Doctrine\ORM\Mapping\Driver\AnnotationDriver;
+
+$config = new Configuration();
+$config->setProxyDir(__DIR__);
+$config->setProxyNamespace('App\GeneratedProxy');
+$config->setMetadataCache(new ArrayCachePool());
+$config->setMetadataDriverImpl(new AnnotationDriver(
+	new AnnotationReader(),
+	[__DIR__ . '/data']
+));
+
+return new EntityManager(
+	DriverManager::getConnection([
+		'driver' => 'mysqli',
+		'memory' => true,
+	]),
+	$config
+);

tests/Type/Doctrine/DBAL/pdo.neon

@@ -0,0 +1,6 @@
+includes:
+	- ../../../../extension.neon
+
+parameters:
+	doctrine:
+		objectManagerLoader: pdo.php

tests/Type/Doctrine/DBAL/pdo.php

@@ -0,0 +1,25 @@
+<?php declare(strict_types = 1);
+
+use Cache\Adapter\PHPArray\ArrayCachePool;
+use Doctrine\Common\Annotations\AnnotationReader;
+use Doctrine\DBAL\DriverManager;
+use Doctrine\ORM\Configuration;
+use Doctrine\ORM\EntityManager;
+use Doctrine\ORM\Mapping\Driver\AnnotationDriver;
+
+$config = new Configuration();
+$config->setProxyDir(__DIR__);
+$config->setProxyNamespace('App\GeneratedProxy');
+$config->setMetadataCache(new ArrayCachePool());
+$config->setMetadataDriverImpl(new AnnotationDriver(
+	new AnnotationReader(),
+	[__DIR__ . '/data']
+));
+
+return new EntityManager(
+	DriverManager::getConnection([
+		'driver' => 'pdo_pgsql',
+		'memory' => true,
+	]),
+	$config
+);