Improve QueryResultDynamicReturnTypeExtension

Author: VincentLanglet

src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php

@@ -16,8 +16,10 @@
 use PHPStan\Type\IterableType;
 use PHPStan\Type\MixedType;
 use PHPStan\Type\NullType;
+use PHPStan\Type\ObjectWithoutClassType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\TypeTraverser;
 use PHPStan\Type\VoidType;
 
 final class QueryResultDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
@@ -109,12 +111,20 @@ private function getMethodReturnTypeForHydrationMode(
 			return $this->originalReturnType($methodReflection);
 		}
 
-		if (!$this->isObjectHydrationMode($hydrationMode)) {
-			// We support only HYDRATE_OBJECT. For other hydration modes, we
-			// return the declared return type of the method.
+		if (!$hydrationMode instanceof ConstantIntegerType) {
 			return $this->originalReturnType($methodReflection);
 		}
 
+		switch ($hydrationMode->getValue()) {
+			case AbstractQuery::HYDRATE_OBJECT;
+				break;
+			case AbstractQuery::HYDRATE_ARRAY;
+				$queryResultType = $this->getArrayHydratedReturnType($queryResultType);
+				break;
+			default:
+				return $this->originalReturnType($methodReflection);
+		}
+
 		switch ($methodReflection->getName()) {
 			case 'getSingleResult':
 				return $queryResultType;
@@ -133,13 +143,22 @@ private function getMethodReturnTypeForHydrationMode(
 		}
 	}
 
-	private function isObjectHydrationMode(Type $type): bool
+	private function getArrayHydratedReturnType(Type $queryResultType): Type
 	{
-		if (!$type instanceof ConstantIntegerType) {
-			return false;
-		}
-
-		return $type->getValue() === AbstractQuery::HYDRATE_OBJECT;
+		return TypeTraverser::map(
+			$queryResultType,
+			static function (Type $type, callable $traverse): Type {
+				$isObject = (new ObjectWithoutClassType())->isSuperTypeOf($type);
+				if ($isObject->yes()) {
+					return new ArrayType(new MixedType(), new MixedType());
+				}
+				if ($isObject->maybe()) {
+					return new MixedType();
+				}
+
+				return $traverse($type);
+			}
+		);
 	}
 
 	private function originalReturnType(MethodReflection $methodReflection): Type