Repository: phpstan/phpstan-doctrine (php)

Issue:
Keep QueryBuilderType across method calls
I don't  know if this is currently even possible, but it would be nice, if Phpstan could determine the type of the QueryBuilder, even if it is defined in another method. Currently something like this results in a mixed return type:
```php
    /**
     * @return array<Entity>
     */
    public function reuseQueryBuilder(): array
    {
        return $this->getQueryBuilder()
            ->getQuery()
            ->getResult();
    }

    private function getQueryBuilder(): QueryBuilder
    {
        return $this->createQueryBuilder('test');
    }
```

Pull request description:
Determine queryBuilderType if getQuery is called immediately
Closes #266 

Uses `searchOtherMethodsForQueryBuilderBeginning` also in `QueryBuilderGetQueryDynamicReturnTypeExtension`.

Patch:
diff --git a/README.md b/README.md
index 228c132..994c509 100644
--- a/README.md
+++ b/README.md
@@ -17,7 +17,7 @@ This extension provides following features:
 * Provides correct return type for `Doctrine\ORM\EntityManager::find`, `getReference` and `getPartialReference` when `Foo::class` entity class name is provided as the first argument
 * Adds missing `matching` method on `Doctrine\Common\Collections\Collection`. This can be turned off by setting `parameters.doctrine.allCollectionsSelectable` to `false`.
 * Also supports Doctrine ODM.
-* Analysis of discrepancies between entity column types and property field types.
+* Analysis of discrepancies between entity column types and property field types. This can be relaxed with the `allowNullablePropertyForRequiredField: true` setting.
 
 ## Installation
 
diff --git a/extension.neon b/extension.neon
index 24e7f9f..f64962d 100644
--- a/extension.neon
+++ b/extension.neon
@@ -110,9 +110,7 @@ services:
 	-
 		class: PHPStan\Type\Doctrine\QueryBuilder\QueryBuilderMethodDynamicReturnTypeExtension
 		arguments:
-			parser: @defaultAnalysisParser
 			queryBuilderClass: %doctrine.queryBuilderClass%
-			descendIntoOtherMethods: %doctrine.searchOtherMethodsForQueryBuilderBeginning%
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
 
@@ -147,6 +145,11 @@ services:
 		class: PHPStan\Rules\Doctrine\ORM\PropertiesExtension
 		tags:
 			- phpstan.properties.readWriteExtension
+	-
+		class: PHPStan\Type\Doctrine\QueryBuilder\OtherMethodQueryBuilderParser
+		arguments:
+			descendIntoOtherMethods: %doctrine.searchOtherMethodsForQueryBuilderBeginning%
+			parser: @defaultAnalysisParser
 
 	doctrineQueryBuilderArgumentsProcessor:
 		class: PHPStan\Type\Doctrine\ArgumentsProcessor
diff --git a/phpstan.neon b/phpstan.neon
index e1d6a3c..93aed16 100644
--- a/phpstan.neon
+++ b/phpstan.neon
@@ -26,4 +26,4 @@ parameters:
 			path: src/Type/Doctrine/QueryBuilder/Expr/ExpressionBuilderDynamicReturnTypeExtension.php
 		-
 			message: '~^Variable property access on PhpParser\\Node\\Stmt\\Declare_\|PhpParser\\Node\\Stmt\\Namespace_\.$~'
-			path: src/Type/Doctrine/QueryBuilder/QueryBuilderMethodDynamicReturnTypeExtension.php
+			path: src/Type/Doctrine/QueryBuilder/OtherMethodQueryBuilderParser.php
diff --git a/src/Type/Doctrine/QueryBuilder/OtherMethodQueryBuilderParser.php b/src/Type/Doctrine/QueryBuilder/OtherMethodQueryBuilderParser.php
new file mode 100644
index 0000000..40116c1
--- /dev/null
+++ b/src/Type/Doctrine/QueryBuilder/OtherMethodQueryBuilderParser.php
@@ -0,0 +1,188 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\QueryBuilder;
+
+use PhpParser\Node;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Identifier;
+use PhpParser\Node\Stmt;
+use PhpParser\Node\Stmt\Class_;
+use PhpParser\Node\Stmt\ClassMethod;
+use PhpParser\Node\Stmt\Declare_;
+use PhpParser\Node\Stmt\Namespace_;
+use PhpParser\Node\Stmt\Return_;
+use PHPStan\Analyser\NodeScopeResolver;
+use PHPStan\Analyser\Scope;
+use PHPStan\Analyser\ScopeContext;
+use PHPStan\Analyser\ScopeFactory;
+use PHPStan\DependencyInjection\Container;
+use PHPStan\Parser\Parser;
+use PHPStan\Reflection\ReflectionProvider;
+use PHPStan\Type\Generic\TemplateTypeMap;
+use PHPStan\Type\TypeWithClassName;
+use function is_array;
+
+class OtherMethodQueryBuilderParser
+{
+
+	/** @var bool */
+	private $descendIntoOtherMethods;
+
+	/** @var ReflectionProvider */
+	private $broker;
+
+	/** @var Parser */
+	private $parser;
+
+	/** @var Container */
+	private $container;
+
+	public function __construct(bool $descendIntoOtherMethods, ReflectionProvider $broker, Parser $parser, Container $container)
+	{
+		$this->descendIntoOtherMethods = $descendIntoOtherMethods;
+		$this->broker = $broker;
+		$this->parser = $parser;
+		$this->container = $container;
+	}
+
+	/**
+	 * @return QueryBuilderType[]
+	 */
+	public function getQueryBuilderTypes(Scope $scope, MethodCall $methodCall): array
+	{
+		if (!$this->descendIntoOtherMethods || !$methodCall->var instanceof MethodCall) {
+			return [];
+		}
+
+		return $this->findQueryBuilderTypesInCalledMethod($scope, $methodCall->var);
+	}
+	/**
+	 * @return QueryBuilderType[]
+	 */
+	private function findQueryBuilderTypesInCalledMethod(Scope $scope, MethodCall $methodCall): array
+	{
+		$methodCalledOnType = $scope->getType($methodCall->var);
+		if (!$methodCall->name instanceof Identifier) {
+			return [];
+		}
+
+		if (!$methodCalledOnType instanceof TypeWithClassName) {
+			return [];
+		}
+
+		if (!$this->broker->hasClass($methodCalledOnType->getClassName())) {
+			return [];
+		}
+
+		$classReflection = $this->broker->getClass($methodCalledOnType->getClassName());
+		$methodName = $methodCall->name->toString();
+		if (!$classReflection->hasNativeMethod($methodName)) {
+			return [];
+		}
+
+		$methodReflection = $classReflection->getNativeMethod($methodName);
+		$fileName = $methodReflection->getDeclaringClass()->getFileName();
+		if ($fileName === null) {
+			return [];
+		}
+
+		$nodes = $this->parser->parseFile($fileName);
+		$classNode = $this->findClassNode($methodReflection->getDeclaringClass()->getName(), $nodes);
+		if ($classNode === null) {
+			return [];
+		}
+
+		$methodNode = $this->findMethodNode($methodReflection->getName(), $classNode->stmts);
+		if ($methodNode === null || $methodNode->stmts === null) {
+			return [];
+		}
+
+		/** @var NodeScopeResolver $nodeScopeResolver */
+		$nodeScopeResolver = $this->container->getByType(NodeScopeResolver::class);
+
+		/** @var ScopeFactory $scopeFactory */
+		$scopeFactory = $this->container->getByType(ScopeFactory::class);
+
+		$methodScope = $scopeFactory->create(
+			ScopeContext::create($fileName),
+			$scope->isDeclareStrictTypes(),
+			[],
+			$methodReflection,
+			$scope->getNamespace()
+		)->enterClass($methodReflection->getDeclaringClass())->enterClassMethod($methodNode, TemplateTypeMap::createEmpty(), [], null, null, null, false, false, false);
+
+		$queryBuilderTypes = [];
+
+		$nodeScopeResolver->processNodes($methodNode->stmts, $methodScope, static function (Node $node, Scope $scope) use (&$queryBuilderTypes): void {
+			if (!$node instanceof Return_ || $node->expr === null) {
+				return;
+			}
+
+			$exprType = $scope->getType($node->expr);
+			if (!$exprType instanceof QueryBuilderType) {
+				return;
+			}
+
+			$queryBuilderTypes[] = $exprType;
+		});
+
+		return $queryBuilderTypes;
+	}
+
+	/**
+	 * @param Node[] $nodes
+	 */
+	private function findClassNode(string $className, array $nodes): ?Class_
+	{
+		foreach ($nodes as $node) {
+			if (
+				$node instanceof Class_
+				&& $node->namespacedName !== null
+				&& $node->namespacedName->toString() === $className
+			) {
+				return $node;
+			}
+
+			if (
+				!$node instanceof Namespace_
+				&& !$node instanceof Declare_
+			) {
+				continue;
+			}
+			$subNodeNames = $node->getSubNodeNames();
+			foreach ($subNodeNames as $subNodeName) {
+				$subNode = $node->{$subNodeName};
+				if (!is_array($subNode)) {
+					$subNode = [$subNode];
+				}
+
+				$result = $this->findClassNode($className, $subNode);
+				if ($result === null) {
+					continue;
+				}
+
+				return $result;
+			}
+		}
+
+		return null;
+	}
+
+	/**
+	 * @param Stmt[] $classStatements
+	 */
+	private function findMethodNode(string $methodName, array $classStatements): ?ClassMethod
+	{
+		foreach ($classStatements as $statement) {
+			if (
+				$statement instanceof ClassMethod
+				&& $statement->name->toString() === $methodName
+			) {
+				return $statement;
+			}
+		}
+
+		return null;
+	}
+
+}
diff --git a/src/Type/Doctrine/QueryBuilder/QueryBuilderGetQueryDynamicReturnTypeExtension.php b/src/Type/Doctrine/QueryBuilder/QueryBuilderGetQueryDynamicReturnTypeExtension.php
index a2a09e7..c828411 100644
--- a/src/Type/Doctrine/QueryBuilder/QueryBuilderGetQueryDynamicReturnTypeExtension.php
+++ b/src/Type/Doctrine/QueryBuilder/QueryBuilderGetQueryDynamicReturnTypeExtension.php
@@ -42,17 +42,22 @@ class QueryBuilderGetQueryDynamicReturnTypeExtension implements DynamicMethodRet
 	/** @var DescriptorRegistry */
 	private $descriptorRegistry;
 
+	/** @var OtherMethodQueryBuilderParser */
+	private $otherMethodQueryBuilderParser;
+
 	public function __construct(
 		ObjectMetadataResolver $objectMetadataResolver,
 		ArgumentsProcessor $argumentsProcessor,
 		?string $queryBuilderClass,
-		DescriptorRegistry $descriptorRegistry
+		DescriptorRegistry $descriptorRegistry,
+		OtherMethodQueryBuilderParser $otherMethodQueryBuilderParser
 	)
 	{
 		$this->objectMetadataResolver = $objectMetadataResolver;
 		$this->argumentsProcessor = $argumentsProcessor;
 		$this->queryBuilderClass = $queryBuilderClass;
 		$this->descriptorRegistry = $descriptorRegistry;
+		$this->otherMethodQueryBuilderParser = $otherMethodQueryBuilderParser;
 	}
 
 	public function getClass(): string
@@ -79,7 +84,10 @@ class QueryBuilderGetQueryDynamicReturnTypeExtension implements DynamicMethodRet
 		)->getReturnType();
 		$queryBuilderTypes = DoctrineTypeUtils::getQueryBuilderTypes($calledOnType);
 		if (count($queryBuilderTypes) === 0) {
-			return $defaultReturnType;
+			$queryBuilderTypes = $this->otherMethodQueryBuilderParser->getQueryBuilderTypes($scope, $methodCall);
+			if (count($queryBuilderTypes) === 0) {
+				return $defaultReturnType;
+			}
 		}
 
 		$objectManager = $this->objectMetadataResolver->getObjectManager();
diff --git a/src/Type/Doctrine/QueryBuilder/QueryBuilderMethodDynamicReturnTypeExtension.php b/src/Type/Doctrine/QueryBuilder/QueryBuilderMethodDynamicReturnTypeExtension.php
index 449e35c..32b437f 100644
--- a/src/Type/Doctrine/QueryBuilder/QueryBuilderMethodDynamicReturnTypeExtension.php
+++ b/src/Type/Doctrine/QueryBuilder/QueryBuilderMethodDynamicReturnTypeExtension.php
@@ -3,74 +3,39 @@
 namespace PHPStan\Type\Doctrine\QueryBuilder;
 
 use Doctrine\ORM\QueryBuilder;
-use PhpParser\Node;
 use PhpParser\Node\Expr\MethodCall;
 use PhpParser\Node\Identifier;
-use PhpParser\Node\Stmt;
-use PhpParser\Node\Stmt\Class_;
-use PhpParser\Node\Stmt\ClassMethod;
-use PhpParser\Node\Stmt\Declare_;
-use PhpParser\Node\Stmt\Namespace_;
-use PhpParser\Node\Stmt\Return_;
-use PHPStan\Analyser\NodeScopeResolver;
 use PHPStan\Analyser\Scope;
-use PHPStan\Analyser\ScopeContext;
-use PHPStan\Analyser\ScopeFactory;
-use PHPStan\Broker\Broker;
-use PHPStan\DependencyInjection\Container;
-use PHPStan\Parser\Parser;
-use PHPStan\Reflection\BrokerAwareExtension;
 use PHPStan\Reflection\MethodReflection;
 use PHPStan\Reflection\ParametersAcceptorSelector;
 use PHPStan\Type\Doctrine\DoctrineTypeUtils;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
-use PHPStan\Type\Generic\TemplateTypeMap;
 use PHPStan\Type\MixedType;
 use PHPStan\Type\ObjectType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
-use PHPStan\Type\TypeWithClassName;
 use function count;
 use function in_array;
-use function is_array;
 use function strtolower;
 
-class QueryBuilderMethodDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension, BrokerAwareExtension
+class QueryBuilderMethodDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
 {
 
 	private const MAX_COMBINATIONS = 16;
 
-	/** @var Container */
-	private $container;
-
-	/** @var Parser */
-	private $parser;
-
 	/** @var string|null */
 	private $queryBuilderClass;
 
-	/** @var bool */
-	private $descendIntoOtherMethods;
-
-	/** @var Broker */
-	private $broker;
+	/** @var OtherMethodQueryBuilderParser */
+	private $otherMethodQueryBuilderParser;
 
 	public function __construct(
-		Container $container,
-		Parser $parser,
 		?string $queryBuilderClass,
-		bool $descendIntoOtherMethods
+		OtherMethodQueryBuilderParser $otherMethodQueryBuilderParser
 	)
 	{
-		$this->container = $container;
-		$this->parser = $parser;
 		$this->queryBuilderClass = $queryBuilderClass;
-		$this->descendIntoOtherMethods = $descendIntoOtherMethods;
-	}
-
-	public function setBroker(Broker $broker): void
-	{
-		$this->broker = $broker;
+		$this->otherMethodQueryBuilderParser = $otherMethodQueryBuilderParser;
 	}
 
 	public function getClass(): string
@@ -109,11 +74,7 @@ class QueryBuilderMethodDynamicReturnTypeExtension implements DynamicMethodRetur
 
 		$queryBuilderTypes = DoctrineTypeUtils::getQueryBuilderTypes($calledOnType);
 		if (count($queryBuilderTypes) === 0) {
-			if (!$this->descendIntoOtherMethods || !$methodCall->var instanceof MethodCall) {
-				return $calledOnType;
-			}
-
-			$queryBuilderTypes = $this->findQueryBuilderTypesInCalledMethod($scope, $methodCall->var);
+			$queryBuilderTypes = $this->otherMethodQueryBuilderParser->getQueryBuilderTypes($scope, $methodCall);
 			if (count($queryBuilderTypes) === 0) {
 				return $calledOnType;
 			}
@@ -131,133 +92,4 @@ class QueryBuilderMethodDynamicReturnTypeExtension implements DynamicMethodRetur
 		return TypeCombinator::union(...$resultTypes);
 	}
 
-	/**
-	 * @return QueryBuilderType[]
-	 */
-	private function findQueryBuilderTypesInCalledMethod(Scope $scope, MethodCall $methodCall): array
-	{
-		$methodCalledOnType = $scope->getType($methodCall->var);
-		if (!$methodCall->name instanceof Identifier) {
-			return [];
-		}
-
-		if (!$methodCalledOnType instanceof TypeWithClassName) {
-			return [];
-		}
-
-		if (!$this->broker->hasClass($methodCalledOnType->getClassName())) {
-			return [];
-		}
-
-		$classReflection = $this->broker->getClass($methodCalledOnType->getClassName());
-		$methodName = $methodCall->name->toString();
-		if (!$classReflection->hasNativeMethod($methodName)) {
-			return [];
-		}
-
-		$methodReflection = $classReflection->getNativeMethod($methodName);
-		$fileName = $methodReflection->getDeclaringClass()->getFileName();
-		if ($fileName === null) {
-			return [];
-		}
-
-		$nodes = $this->parser->parseFile($fileName);
-		$classNode = $this->findClassNode($methodReflection->getDeclaringClass()->getName(), $nodes);
-		if ($classNode === null) {
-			return [];
-		}
-
-		$methodNode = $this->findMethodNode($methodReflection->getName(), $classNode->stmts);
-		if ($methodNode === null || $methodNode->stmts === null) {
-			return [];
-		}
-
-		/** @var NodeScopeResolver $nodeScopeResolver */
-		$nodeScopeResolver = $this->container->getByType(NodeScopeResolver::class);
-
-		/** @var ScopeFactory $scopeFactory */
-		$scopeFactory = $this->container->getByType(ScopeFactory::class);
-
-		$methodScope = $scopeFactory->create(
-			ScopeContext::create($fileName),
-			$scope->isDeclareStrictTypes(),
-			[],
-			$methodReflection,
-			$scope->getNamespace()
-		)->enterClass($methodReflection->getDeclaringClass())->enterClassMethod($methodNode, TemplateTypeMap::createEmpty(), [], null, null, null, false, false, false);
-
-		$queryBuilderTypes = [];
-
-		$nodeScopeResolver->processNodes($methodNode->stmts, $methodScope, static function (Node $node, Scope $scope) use (&$queryBuilderTypes): void {
-			if (!$node instanceof Return_ || $node->expr === null) {
-				return;
-			}
-
-			$exprType = $scope->getType($node->expr);
-			if (!$exprType instanceof QueryBuilderType) {
-				return;
-			}
-
-			$queryBuilderTypes[] = $exprType;
-		});
-
-		return $queryBuilderTypes;
-	}
-
-	/**
-	 * @param Node[] $nodes
-	 */
-	private function findClassNode(string $className, array $nodes): ?Class_
-	{
-		foreach ($nodes as $node) {
-			if (
-				$node instanceof Class_
-				&& $node->namespacedName !== null
-				&& $node->namespacedName->toString() === $className
-			) {
-				return $node;
-			}
-
-			if (
-				!$node instanceof Namespace_
-				&& !$node instanceof Declare_
-			) {
-				continue;
-			}
-			$subNodeNames = $node->getSubNodeNames();
-			foreach ($subNodeNames as $subNodeName) {
-				$subNode = $node->{$subNodeName};
-				if (!is_array($subNode)) {
-					$subNode = [$subNode];
-				}
-
-				$result = $this->findClassNode($className, $subNode);
-				if ($result === null) {
-					continue;
-				}
-
-				return $result;
-			}
-		}
-
-		return null;
-	}
-
-	/**
-	 * @param Stmt[] $classStatements
-	 */
-	private function findMethodNode(string $methodName, array $classStatements): ?ClassMethod
-	{
-		foreach ($classStatements as $statement) {
-			if (
-				$statement instanceof ClassMethod
-				&& $statement->name->toString() === $methodName
-			) {
-				return $statement;
-			}
-		}
-
-		return null;
-	}
-
 }


Test patch:
diff --git a/tests/Type/Doctrine/data/QueryResult/queryBuilderGetQuery.php b/tests/Type/Doctrine/data/QueryResult/queryBuilderGetQuery.php
index a9b47e9..5e9ddd3 100644
--- a/tests/Type/Doctrine/data/QueryResult/queryBuilderGetQuery.php
+++ b/tests/Type/Doctrine/data/QueryResult/queryBuilderGetQuery.php
@@ -66,4 +66,19 @@ class QueryBuilderGetQuery
 		assertType('Doctrine\ORM\Query<void>', $query);
 
 	}
+
+	public function testQueryTypeIsInferredOnAcrossMethods(EntityManagerInterface $em): void
+	{
+		$query = $this->getQueryBuilder($em)
+			->getQuery();
+
+		assertType('Doctrine\ORM\Query<QueryResult\Entities\Many>', $query);
+	}
+
+	private function getQueryBuilder(EntityManagerInterface $em): QueryBuilder
+	{
+		return $em->createQueryBuilder()
+			->select('m')
+			->from(Many::class, 'm');
+	}
 }