Infer QueryBuilderType for any method returning QB

Author: janedbal

extension.neon

@@ -136,6 +136,10 @@ services:
 		class: PHPStan\Type\Doctrine\Query\QueryGetDqlDynamicReturnTypeExtension
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: PHPStan\Type\Doctrine\QueryBuilder\QueryBuilderReturnedDynamicReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension
 	-
 		class: PHPStan\Type\Doctrine\CreateQueryDynamicReturnTypeExtension
 		arguments:

src/Type/Doctrine/QueryBuilder/OtherMethodQueryBuilderParser.php

@@ -54,17 +54,22 @@ public function __construct(bool $descendIntoOtherMethods, ReflectionProvider $r
 	 */
 	public function getQueryBuilderTypes(Scope $scope, MethodCall $methodCall): array
 	{
-		if (!$this->descendIntoOtherMethods || !$methodCall->var instanceof MethodCall) {
+		if (!$methodCall->var instanceof MethodCall) {
 			return [];
 		}
 
 		return $this->findQueryBuilderTypesInCalledMethod($scope, $methodCall->var);
 	}
+
 	/**
 	 * @return QueryBuilderType[]
 	 */
-	private function findQueryBuilderTypesInCalledMethod(Scope $scope, MethodCall $methodCall): array
+	public function findQueryBuilderTypesInCalledMethod(Scope $scope, MethodCall $methodCall): array
 	{
+		if (!$this->descendIntoOtherMethods) {
+			return [];
+		}
+
 		$methodCalledOnType = $scope->getType($methodCall->var);
 		if (!$methodCall->name instanceof Identifier) {
 			return [];

src/Type/Doctrine/QueryBuilder/QueryBuilderReturnedDynamicReturnTypeExtension.php

@@ -0,0 +1,68 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\QueryBuilder;
+
+use Doctrine\ORM\QueryBuilder;
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\MixedType;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use function count;
+
+class QueryBuilderReturnedDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+
+	/** @var OtherMethodQueryBuilderParser */
+	private $otherMethodQueryBuilderParser;
+
+	public function __construct(
+		OtherMethodQueryBuilderParser $otherMethodQueryBuilderParser
+	)
+	{
+		$this->otherMethodQueryBuilderParser = $otherMethodQueryBuilderParser;
+	}
+
+	public function getClass(): string
+	{
+		return 'QueryResult\CreateQuery\QueryBuilderGetQuery'; // TODO https://github.com/phpstan/phpstan-src/pull/2761
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		$returnType = ParametersAcceptorSelector::selectSingle(
+			$methodReflection->getVariants()
+		)->getReturnType();
+
+		if ($returnType instanceof MixedType) {
+			return false;
+		}
+
+		return (new ObjectType(QueryBuilder::class))->isSuperTypeOf($returnType)->yes();
+	}
+
+	public function getTypeFromMethodCall(
+		MethodReflection $methodReflection,
+		MethodCall $methodCall,
+		Scope $scope
+	): Type
+	{
+		$defaultReturnType = ParametersAcceptorSelector::selectFromArgs(
+			$scope,
+			$methodCall->getArgs(),
+			$methodReflection->getVariants()
+		)->getReturnType();
+
+		$queryBuilderTypes = $this->otherMethodQueryBuilderParser->findQueryBuilderTypesInCalledMethod($scope, $methodCall);
+		if (count($queryBuilderTypes) === 0) {
+			return $defaultReturnType;
+		}
+
+		return TypeCombinator::union(...$queryBuilderTypes);
+	}
+
+}

tests/Type/Doctrine/data/QueryResult/queryBuilderGetQuery.php

@@ -242,6 +242,13 @@ public function testQueryTypeIsInferredOnAcrossMethods(EntityManagerInterface $e
 		assertType('Doctrine\ORM\Query<null, QueryResult\Entities\Many>', $branchingQuery);
 	}
 
+	public function testQueryTypeIsInferredOnAcrossMethodsEvenWhenVariableAssignmentIsUsed(EntityManagerInterface $em): void
+	{
+		$queryBuilder = $this->getQueryBuilder($em);
+
+		assertType('Doctrine\ORM\Query<null, QueryResult\Entities\Many>', $queryBuilder->getQuery());
+	}
+
 	private function getQueryBuilder(EntityManagerInterface $em): QueryBuilder
 	{
 		return $em->createQueryBuilder()