From 0d2be9c295eda87a14aa979ff1b2289f364c53ee Mon Sep 17 00:00:00 2001
From: Vincent Langlet <vincentlanglet@hotmail.fr>
Date: Wed, 5 Apr 2023 11:50:04 +0200
Subject: [PATCH] Improve return type of getArrayResult

---
 .../QueryResultDynamicReturnTypeExtension.php | 44 ++++++++++++++++---
 .../Doctrine/data/QueryResult/queryResult.php | 16 +++----
 2 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php b/src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php
index e7357b35..8dcb82a8 100644
--- a/src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php
+++ b/src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php
@@ -10,7 +10,11 @@
 use PHPStan\ShouldNotHappenException;
 use PHPStan\Type\Accessory\AccessoryArrayListType;
 use PHPStan\Type\ArrayType;
+use PHPStan\Type\Constant\ConstantArrayType;
 use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\Doctrine\DescriptorNotRegisteredException;
+use PHPStan\Type\Doctrine\DescriptorRegistry;
 use PHPStan\Type\Doctrine\ObjectMetadataResolver;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
 use PHPStan\Type\IntegerType;
@@ -48,11 +52,16 @@ final class QueryResultDynamicReturnTypeExtension implements DynamicMethodReturn
 	/** @var ObjectMetadataResolver */
 	private $objectMetadataResolver;
 
+	/** @var DescriptorRegistry */
+	private $descriptorRegistry;
+
 	public function __construct(
-		ObjectMetadataResolver $objectMetadataResolver
+		ObjectMetadataResolver $objectMetadataResolver,
+		DescriptorRegistry $descriptorRegistry
 	)
 	{
 		$this->objectMetadataResolver = $objectMetadataResolver;
+		$this->descriptorRegistry = $descriptorRegistry;
 	}
 
 	public function getClass(): string
@@ -183,10 +192,11 @@ private function getMethodReturnTypeForHydrationMode(
 	private function getArrayHydratedReturnType(Type $queryResultType): Type
 	{
 		$objectManager = $this->objectMetadataResolver->getObjectManager();
+		$descriptorRegistry = $this->descriptorRegistry;
 
 		return TypeTraverser::map(
 			$queryResultType,
-			static function (Type $type, callable $traverse) use ($objectManager): Type {
+			static function (Type $type, callable $traverse) use ($objectManager, $descriptorRegistry): Type {
 				$isObject = (new ObjectWithoutClassType())->isSuperTypeOf($type);
 				if ($isObject->no()) {
 					return $traverse($type);
@@ -199,9 +209,33 @@ static function (Type $type, callable $traverse) use ($objectManager): Type {
 					return new MixedType();
 				}
 
-				return $objectManager->getMetadataFactory()->hasMetadataFor($type->getClassName())
-					? new ArrayType(new MixedType(), new MixedType())
-					: $traverse($type);
+				if (!$objectManager->getMetadataFactory()->hasMetadataFor($type->getClassName())) {
+					return $traverse($type);
+				}
+
+				$metadata = $objectManager->getMetadataFactory()->getMetadataFor($type->getClassName());
+
+				$types = [];
+				$keys = [];
+				foreach ($metadata->fieldMappings as $fieldMapping) {
+					try {
+						$type = $descriptorRegistry->get($fieldMapping['type'])->getWritableToPropertyType();
+					} catch (DescriptorNotRegisteredException $exception) {
+						return new ArrayType(new MixedType(), new MixedType());
+					}
+
+					$nullable = isset($fieldMapping['nullable'])
+						? $fieldMapping['nullable'] === true
+						: false;
+					if ($nullable) {
+						$type = TypeCombinator::addNull($type);
+					}
+
+					$types[] = $type;
+					$keys[] = new ConstantStringType($fieldMapping['fieldName']);
+				}
+
+				return new ConstantArrayType($keys, $types);
 			}
 		);
 	}
diff --git a/tests/Type/Doctrine/data/QueryResult/queryResult.php b/tests/Type/Doctrine/data/QueryResult/queryResult.php
index 61ed83d1..b886e152 100644
--- a/tests/Type/Doctrine/data/QueryResult/queryResult.php
+++ b/tests/Type/Doctrine/data/QueryResult/queryResult.php
@@ -155,35 +155,35 @@ public function testReturnTypeOfQueryMethodsWithExplicitArrayHydrationMode(Entit
 		');
 
 		assertType(
-			'list<array>',
+			'list<array{id: numeric-string, intColumn: int, stringColumn: string, stringNullColumn: string|null, datetimeColumn: DateTime, datetimeImmutableColumn: DateTimeImmutable}>',
 			$query->getResult(AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'list<array>',
+			'list<array{id: numeric-string, intColumn: int, stringColumn: string, stringNullColumn: string|null, datetimeColumn: DateTime, datetimeImmutableColumn: DateTimeImmutable}>',
 			$query->getArrayResult()
 		);
 		assertType(
-			'iterable<int, array>',
+			'iterable<int, array{id: numeric-string, intColumn: int, stringColumn: string, stringNullColumn: string|null, datetimeColumn: DateTime, datetimeImmutableColumn: DateTimeImmutable}>',
 			$query->toIterable([], AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'list<array>',
+			'list<array{id: numeric-string, intColumn: int, stringColumn: string, stringNullColumn: string|null, datetimeColumn: DateTime, datetimeImmutableColumn: DateTimeImmutable}>',
 			$query->execute(null, AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'list<array>',
+			'list<array{id: numeric-string, intColumn: int, stringColumn: string, stringNullColumn: string|null, datetimeColumn: DateTime, datetimeImmutableColumn: DateTimeImmutable}>',
 			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'list<array>',
+			'list<array{id: numeric-string, intColumn: int, stringColumn: string, stringNullColumn: string|null, datetimeColumn: DateTime, datetimeImmutableColumn: DateTimeImmutable}>',
 			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'array',
+			'array{id: numeric-string, intColumn: int, stringColumn: string, stringNullColumn: string|null, datetimeColumn: DateTime, datetimeImmutableColumn: DateTimeImmutable}',
 			$query->getSingleResult(AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'array|null',
+			'array{id: numeric-string, intColumn: int, stringColumn: string, stringNullColumn: string|null, datetimeColumn: DateTime, datetimeImmutableColumn: DateTimeImmutable}|null',
 			$query->getOneOrNullResult(AbstractQuery::HYDRATE_ARRAY)
 		);