Add DynamicReturnTypeExtension for Query::getResult() and variants

Author: arnaud-lb

extension.neon

@@ -127,6 +127,10 @@ services:
 			objectMetadataResolver: @PHPStan\Type\Doctrine\ObjectMetadataResolver
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: PHPStan\Type\Doctrine\Query\QueryResultDynamicReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension
 	-
 		class: PHPStan\Type\Doctrine\QueryBuilder\Expr\ExpressionBuilderDynamicReturnTypeExtension
 		arguments:

src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php

@@ -0,0 +1,139 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\Query;
+
+use Doctrine\ORM\AbstractQuery;
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\ShouldNotHappenException;
+use PHPStan\Type\ArrayType;
+use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\Generic\GenericObjectType;
+use PHPStan\Type\IntegerType;
+use PHPStan\Type\MixedType;
+use PHPStan\Type\NullType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\VoidType;
+
+final class QueryResultDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+
+	private const METHOD_HYDRATION_MODE_ARG = [
+		'getResult' => 0,
+		'execute' => 1,
+		'executeIgnoreQueryCache' => 1,
+		'executeUsingQueryCache' => 1,
+		'getOneOrNullResult' => 0,
+		'getSingleResult' => 0,
+	];
+
+	public function getClass(): string
+	{
+		return AbstractQuery::class;
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		return isset(self::METHOD_HYDRATION_MODE_ARG[$methodReflection->getName()]);
+	}
+
+	public function getTypeFromMethodCall(
+		MethodReflection $methodReflection,
+		MethodCall $methodCall,
+		Scope $scope
+	): Type
+	{
+		$methodName = $methodReflection->getName();
+
+		if (!isset(self::METHOD_HYDRATION_MODE_ARG[$methodName])) {
+			throw new ShouldNotHappenException();
+		}
+
+		$argIndex = self::METHOD_HYDRATION_MODE_ARG[$methodName];
+		$args = $methodCall->getArgs();
+
+		if (isset($args[$argIndex])) {
+			$hydrationMode = $scope->getType($args[$argIndex]->value);
+		} else {
+			$parametersAcceptor = ParametersAcceptorSelector::selectSingle(
+				$methodReflection->getVariants()
+			);
+			$parameter = $parametersAcceptor->getParameters()[$argIndex];
+			$hydrationMode = $parameter->getDefaultValue() ?? new NullType();
+		}
+
+		$queryType = $scope->getType($methodCall->var);
+		$queryResultType = $this->getQueryResultType($queryType);
+
+		return $this->getMethodReturnTypeForHydrationMode(
+			$methodReflection,
+			$hydrationMode,
+			$queryResultType
+		);
+	}
+
+	private function getQueryResultType(Type $queryType): Type
+	{
+		if (!$queryType instanceof GenericObjectType) {
+			return new MixedType();
+		}
+
+		$types = $queryType->getTypes();
+
+		return $types[0] ?? new MixedType();
+	}
+
+	private function getMethodReturnTypeForHydrationMode(
+		MethodReflection $methodReflection,
+		Type $hydrationMode,
+		Type $queryResultType
+	): Type
+	{
+		if ($queryResultType instanceof VoidType) {
+			// A void query result type indicates an UPDATE or DELETE query.
+			// In this case all methods return the number of affected rows.
+			return new IntegerType();
+		}
+
+		if (!$this->isObjectHydrationMode($hydrationMode)) {
+			// We support only HYDRATE_OBJECT. For other hydration modes, we
+			// return the declared return type of the method.
+			return $this->originalReturnType($methodReflection);
+		}
+
+		switch ($methodReflection->getName()) {
+			case 'getSingleResult':
+				return $queryResultType;
+			case 'getOneOrNullResult':
+				return TypeCombinator::addNull($queryResultType);
+			default:
+				return new ArrayType(
+					new MixedType(),
+					$queryResultType
+				);
+		}
+	}
+
+	private function isObjectHydrationMode(Type $type): bool
+	{
+		if (!$type instanceof ConstantIntegerType) {
+			return false;
+		}
+
+		return $type->getValue() === AbstractQuery::HYDRATE_OBJECT;
+	}
+
+	private function originalReturnType(MethodReflection $methodReflection): Type
+	{
+		$parametersAcceptor = ParametersAcceptorSelector::selectSingle(
+			$methodReflection->getVariants()
+		);
+
+		return $parametersAcceptor->getReturnType();
+	}
+
+}

tests/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtensionTest.php

@@ -0,0 +1,37 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\Query;
+
+use PHPStan\Testing\TypeInferenceTestCase;
+
+class QueryResultDynamicReturnTypeExtensionTest extends TypeInferenceTestCase
+{
+
+	/** @return iterable<mixed> */
+	public function dataFileAsserts(): iterable
+	{
+		yield from $this->gatherAssertTypes(__DIR__ . '/../data/QueryResult/queryResult.php');
+	}
+
+	/**
+	 * @dataProvider dataFileAsserts
+	 * @param string $assertType
+	 * @param string $file
+	 * @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__ . '/../data/QueryResult/config.neon'];
+	}
+
+}