Improve QueryResultDynamicReturnTypeExtension

Author: VincentLanglet

src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php

@@ -10,14 +10,22 @@
 use PHPStan\ShouldNotHappenException;
 use PHPStan\Type\Accessory\AccessoryArrayListType;
 use PHPStan\Type\ArrayType;
+use PHPStan\Type\BenevolentUnionType;
 use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\Doctrine\ObjectMetadataResolver;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
 use PHPStan\Type\IntegerType;
 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\TypeUtils;
+use PHPStan\Type\TypeWithClassName;
 use PHPStan\Type\VoidType;
+use function count;
 
 final class QueryResultDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
 {
@@ -32,14 +40,32 @@ final class QueryResultDynamicReturnTypeExtension implements DynamicMethodReturn
 		'getSingleResult' => 0,
 	];
 
+	private const METHOD_HYDRATION_MODE = [
+		'getArrayResult' => AbstractQuery::HYDRATE_ARRAY,
+		'getScalarResult' => AbstractQuery::HYDRATE_SCALAR,
+		'getSingleColumnResult' => AbstractQuery::HYDRATE_SCALAR_COLUMN,
+		'getSingleScalarResult' => AbstractQuery::HYDRATE_SINGLE_SCALAR,
+	];
+
+	/** @var ObjectMetadataResolver */
+	private $objectMetadataResolver;
+
+	public function __construct(
+		ObjectMetadataResolver $objectMetadataResolver
+	)
+	{
+		$this->objectMetadataResolver = $objectMetadataResolver;
+	}
+
 	public function getClass(): string
 	{
 		return AbstractQuery::class;
 	}
 
 	public function isMethodSupported(MethodReflection $methodReflection): bool
 	{
-		return isset(self::METHOD_HYDRATION_MODE_ARG[$methodReflection->getName()]);
+		return isset(self::METHOD_HYDRATION_MODE_ARG[$methodReflection->getName()])
+			|| isset(self::METHOD_HYDRATION_MODE[$methodReflection->getName()]);
 	}
 
 	public function getTypeFromMethodCall(
@@ -50,21 +76,23 @@ public function getTypeFromMethodCall(
 	{
 		$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(self::METHOD_HYDRATION_MODE[$methodName])) {
+			$hydrationMode = new ConstantIntegerType(self::METHOD_HYDRATION_MODE[$methodName]);
+		} elseif (isset(self::METHOD_HYDRATION_MODE_ARG[$methodName])) {
+			$argIndex = self::METHOD_HYDRATION_MODE_ARG[$methodName];
+			$args = $methodCall->getArgs();
 
-		if (isset($args[$argIndex])) {
-			$hydrationMode = $scope->getType($args[$argIndex]->value);
+			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();
+			}
 		} else {
-			$parametersAcceptor = ParametersAcceptorSelector::selectSingle(
-				$methodReflection->getVariants()
-			);
-			$parameter = $parametersAcceptor->getParameters()[$argIndex];
-			$hydrationMode = $parameter->getDefaultValue() ?? new NullType();
+			throw new ShouldNotHappenException();
 		}
 
 		$queryType = $scope->getType($methodCall->var);
@@ -98,23 +126,58 @@ private function getMethodReturnTypeForHydrationMode(
 			return null;
 		}
 
-		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 null;
+		}
+
+		$singleResult = false;
+		switch ($hydrationMode->getValue()) {
+			case AbstractQuery::HYDRATE_OBJECT:
+				break;
+			case AbstractQuery::HYDRATE_ARRAY:
+				$queryResultType = $this->getArrayHydratedReturnType($queryResultType);
+				break;
+			case AbstractQuery::HYDRATE_SCALAR:
+				$queryResultType = $this->getScalarHydratedReturnType($queryResultType);
+				break;
+			case AbstractQuery::HYDRATE_SINGLE_SCALAR:
+				$singleResult = true;
+				$queryResultType = $this->getSingleScalarHydratedReturnType($queryResultType);
+				break;
+			case AbstractQuery::HYDRATE_SIMPLEOBJECT:
+				$queryResultType = $this->getSimpleObjectHydratedReturnType($queryResultType);
+				break;
+			case AbstractQuery::HYDRATE_SCALAR_COLUMN:
+				$queryResultType = $this->getScalarColumnHydratedReturnType($queryResultType);
+				break;
+			default:
+				return null;
+		}
+
+		if ($queryResultType === null) {
 			return null;
 		}
 
 		switch ($methodReflection->getName()) {
 			case 'getSingleResult':
 				return $queryResultType;
 			case 'getOneOrNullResult':
-				return TypeCombinator::addNull($queryResultType);
+				$nullableQueryResultType = TypeCombinator::addNull($queryResultType);
+				if ($queryResultType instanceof BenevolentUnionType) {
+					$nullableQueryResultType = TypeUtils::toBenevolentUnion($nullableQueryResultType);
+				}
+
+				return $nullableQueryResultType;
 			case 'toIterable':
 				return new IterableType(
 					$queryKeyType->isNull()->yes() ? new IntegerType() : $queryKeyType,
 					$queryResultType
 				);
 			default:
+				if ($singleResult) {
+					return $queryResultType;
+				}
+
 				if ($queryKeyType->isNull()->yes()) {
 					return AccessoryArrayListType::intersectWith(new ArrayType(
 						new IntegerType(),
@@ -128,13 +191,127 @@ private function getMethodReturnTypeForHydrationMode(
 		}
 	}
 
-	private function isObjectHydrationMode(Type $type): bool
+	/**
+	 * When we're array-hydrating object, we're not sure of the shape of the array.
+	 * We could return `new ArrayTyp(new MixedType(), new MixedType())`
+	 * but the lack of precision in the array keys/values would give false positive.
+	 *
+	 * @see https://github.com/phpstan/phpstan-doctrine/pull/412#issuecomment-1497092934
+	 */
+	private function getArrayHydratedReturnType(Type $queryResultType): ?Type
 	{
-		if (!$type instanceof ConstantIntegerType) {
-			return false;
+		$objectManager = $this->objectMetadataResolver->getObjectManager();
+
+		$mixedFound = false;
+		$queryResultType = TypeTraverser::map(
+			$queryResultType,
+			static function (Type $type, callable $traverse) use ($objectManager, &$mixedFound): Type {
+				$isObject = (new ObjectWithoutClassType())->isSuperTypeOf($type);
+				if ($isObject->no()) {
+					return $traverse($type);
+				}
+				if (
+					$isObject->maybe()
+					|| !$type instanceof TypeWithClassName
+					|| $objectManager === null
+				) {
+					$mixedFound = true;
+
+					return new MixedType();
+				}
+
+				/** @var class-string $className */
+				$className = $type->getClassName();
+				if (!$objectManager->getMetadataFactory()->hasMetadataFor($className)) {
+					return $traverse($type);
+				}
+
+				$mixedFound = true;
+
+				return new MixedType();
+			}
+		);
+
+		return $mixedFound ? null : $queryResultType;
+	}
+
+	/**
+	 * When we're scalar-hydrating object, we're not sure of the shape of the array.
+	 * We could return `new ArrayTyp(new MixedType(), new MixedType())`
+	 * but the lack of precision in the array keys/values would give false positive.
+	 *
+	 * @see https://github.com/phpstan/phpstan-doctrine/pull/453#issuecomment-1895415544
+	 */
+	private function getScalarHydratedReturnType(Type $queryResultType): ?Type
+	{
+		if (!$queryResultType->isArray()->yes()) {
+			return null;
+		}
+
+		foreach ($queryResultType->getArrays() as $arrayType) {
+			$itemType = $arrayType->getItemType();
+
+			if (
+				!(new ObjectWithoutClassType())->isSuperTypeOf($itemType)->no()
+				|| !$itemType->isArray()->no()
+			) {
+				// We could return `new ArrayTyp(new MixedType(), new MixedType())`
+				// but the lack of precision in the array keys/values would give false positive
+				// @see https://github.com/phpstan/phpstan-doctrine/pull/453#issuecomment-1895415544
+				return null;
+			}
+		}
+
+		return $queryResultType;
+	}
+
+	private function getSimpleObjectHydratedReturnType(Type $queryResultType): ?Type
+	{
+		if ((new ObjectWithoutClassType())->isSuperTypeOf($queryResultType)->yes()) {
+			return $queryResultType;
+		}
+
+		return null;
+	}
+
+	private function getSingleScalarHydratedReturnType(Type $queryResultType): ?Type
+	{
+		$queryResultType = $this->getScalarHydratedReturnType($queryResultType);
+		if ($queryResultType === null || !$queryResultType->isConstantArray()->yes()) {
+			return null;
+		}
+
+		$types = [];
+		foreach ($queryResultType->getConstantArrays() as $constantArrayType) {
+			$values = $constantArrayType->getValueTypes();
+			if (count($values) !== 1) {
+				return null;
+			}
+
+			$types[] = $constantArrayType->getFirstIterableValueType();
+		}
+
+		return TypeCombinator::union(...$types);
+	}
+
+	private function getScalarColumnHydratedReturnType(Type $queryResultType): ?Type
+	{
+		$queryResultType = $this->getScalarHydratedReturnType($queryResultType);
+		if ($queryResultType === null || !$queryResultType->isConstantArray()->yes()) {
+			return null;
+		}
+
+		$types = [];
+		foreach ($queryResultType->getConstantArrays() as $constantArrayType) {
+			$values = $constantArrayType->getValueTypes();
+			if (count($values) !== 1) {
+				return null;
+			}
+
+			$types[] = $constantArrayType->getFirstIterableValueType();
 		}
 
-		return $type->getValue() === AbstractQuery::HYDRATE_OBJECT;
+		return TypeCombinator::union(...$types);
 	}
 
 }