Query result type 2

extension.neon

@@ -30,6 +30,7 @@ parameters:
 		- stubs/Collections/Selectable.stub
 		- stubs/ORM/QueryBuilder.stub
 		- stubs/ORM/AbstractQuery.stub
+		- stubs/ORM/Query.stub
 		- stubs/ServiceDocumentRepository.stub
 
 parametersSchema:
@@ -111,6 +112,16 @@ services:
 		class: PHPStan\Type\Doctrine\Query\QueryGetDqlDynamicReturnTypeExtension
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: PHPStan\Type\Doctrine\CreateQueryDynamicReturnTypeExtension
+		arguments:
+			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/CreateQueryDynamicReturnTypeExtension.php

@@ -0,0 +1,95 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine;
+
+use Doctrine\Common\CommonException;
+use Doctrine\DBAL\DBALException;
+use Doctrine\ORM\EntityManagerInterface;
+use Doctrine\ORM\ORMException;
+use Doctrine\ORM\Query;
+use PhpParser\Node\Arg;
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\Doctrine\Query\QueryResultTypeMapping;
+use PHPStan\Type\Doctrine\Query\QueryResultTypeWalker;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\Generic\GenericObjectType;
+use PHPStan\Type\MixedType;
+use PHPStan\Type\Type;
+
+final class CreateQueryDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+
+	/** @var ObjectMetadataResolver */
+	private $objectMetadataResolver;
+
+	/** @var DescriptorRegistry */
+	private $descriptorRegistry;
+
+	public function __construct(ObjectMetadataResolver $objectMetadataResolver, DescriptorRegistry $descriptorRegistry)
+	{
+		$this->objectMetadataResolver = $objectMetadataResolver;
+		$this->descriptorRegistry = $descriptorRegistry;
+	}
+
+	public function getClass(): string
+	{
+		return EntityManagerInterface::class;
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		return $methodReflection->getName() === 'createQuery';
+	}
+
+	public function getTypeFromMethodCall(
+		MethodReflection $methodReflection,
+		MethodCall $methodCall,
+		Scope $scope
+	): Type
+	{
+		$queryStringArgIndex = 0;
+
+		if (!isset($methodCall->args[$queryStringArgIndex]) || !$methodCall->args[$queryStringArgIndex] instanceof Arg) {
+			return $this->fallbackType();
+		}
+
+		$argType = $scope->getType($methodCall->args[$queryStringArgIndex]->value);
+		if (!$argType instanceof ConstantStringType) {
+			return $this->fallbackType();
+		}
+
+		$queryString = $argType->getValue();
+
+		$em = $this->objectMetadataResolver->getObjectManager();
+
+		if (!$em instanceof EntityManagerInterface) {
+			return $this->fallbackType();
+		}
+
+		$typeMapping = new QueryResultTypeMapping();
+
+		try {
+			$query = $em->createQuery($queryString);
+			QueryResultTypeWalker::run($query, $typeMapping, $this->descriptorRegistry);
+		} catch (ORMException | DBALException | CommonException $e) {
+			return $this->fallbackType();
+		}
+
+		return new GenericObjectType(
+			Query::class,
+			[$typeMapping->getResultType()]
+		);
+	}
+
+	private function fallbackType(): GenericObjectType
+	{
+		return new GenericObjectType(
+			Query::class,
+			[new MixedType(true)]
+		);
+	}
+
+}

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();
+	}
+
+}

src/Type/Doctrine/Query/QueryResultTypeMapping.php

@@ -0,0 +1,104 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\Query;
+
+use PHPStan\Type\Constant\ConstantArrayTypeBuilder;
+use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\Type;
+use PHPStan\Type\VoidType;
+
+final class QueryResultTypeMapping
+{
+
+	/** @var bool */
+	private $selectQuery = false;
+
+	/** @var array<array-key,Type> */
+	private $entities = [];
+
+	/** @var array<array-key,Type> */
+	private $scalars = [];
+
+	public function setSelectQuery(): void
+	{
+		$this->selectQuery = true;
+	}
+
+	public function isSelectQuery(): bool
+	{
+		return $this->selectQuery;
+	}
+
+	/**
+	 * @param array-key $alias
+	 */
+	public function addEntity($alias, Type $type): void
+	{
+		$this->entities[$alias] = $type;
+	}
+
+	/**
+	 * @return array<array-key,Type>
+	 */
+	public function getEntities(): array
+	{
+		return $this->entities;
+	}
+
+	/**
+	 * @param array-key $alias
+	 */
+	public function addScalar($alias, Type $type): void
+	{
+		$this->scalars[$alias] = $type;
+	}
+
+	/**
+	 * @return array<int,Type>
+	 */
+	public function getScalars(): array
+	{
+		return $this->scalars;
+	}
+
+	public function getResultType(): Type
+	{
+		if (!$this->selectQuery) {
+			return new VoidType();
+		}
+
+		if (count($this->entities) === 1 && count($this->scalars) === 0) {
+			foreach ($this->entities as $entityType) {
+				return $entityType;
+			}
+		}
+
+		$builder = ConstantArrayTypeBuilder::createEmpty();
+
+		foreach ($this->entities as $alias => $entityType) {
+			$offsetType = $this->resolveOffsetType($alias);
+			$builder->setOffsetValueType($offsetType, $entityType);
+		}
+
+		foreach ($this->scalars as $alias => $scalarType) {
+			$offsetType = $this->resolveOffsetType($alias);
+			$builder->setOffsetValueType($offsetType, $scalarType);
+		}
+
+		return $builder->getArray();
+	}
+
+	/**
+	 * @param array-key $alias
+	 */
+	private function resolveOffsetType($alias): Type
+	{
+		if (is_int($alias)) {
+			return new ConstantIntegerType($alias);
+		}
+
+		return new ConstantStringType($alias);
+	}
+
+}