Look for QueryBuilders in called methods

Author: ondrejmirtes

extension.neon

@@ -4,6 +4,7 @@ parameters:
 		queryBuilderClass: null
 		allCollectionsSelectable: true
 		objectManagerLoader: null
+		searchOtherMethodsForQueryBuilderBeginning: true
 
 		# setting this to false might lead to performance problems
 		# it changes the braching logic - with false, queryBuilders from all branches are analysed separately
@@ -15,6 +16,7 @@ parametersSchema:
 		queryBuilderClass: schema(string(), nullable())
 		allCollectionsSelectable: bool()
 		objectManagerLoader: schema(string(), nullable())
+		searchOtherMethodsForQueryBuilderBeginning: bool()
 		queryBuilderFastAlgorithm: bool()
 	])
 
@@ -74,6 +76,7 @@ services:
 		class: PHPStan\Type\Doctrine\QueryBuilder\QueryBuilderMethodDynamicReturnTypeExtension
 		arguments:
 			queryBuilderClass: %doctrine.queryBuilderClass%
+			descendIntoOtherMethods: %doctrine.searchOtherMethodsForQueryBuilderBeginning%
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
 

phpstan.neon

@@ -18,3 +18,6 @@ parameters:
 		-
 			message: '~^Variable method call on Doctrine\\ORM\\Query\\Expr\.$~'
 			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

rules.neon

@@ -8,6 +8,7 @@ parametersSchema:
 		queryBuilderClass: schema(string(), nullable())
 		allCollectionsSelectable: bool()
 		objectManagerLoader: schema(string(), nullable())
+		searchOtherMethodsForQueryBuilderBeginning: bool()
 		queryBuilderFastAlgorithm: bool()
 		reportDynamicQueryBuilders: bool()
 	])

src/Type/Doctrine/QueryBuilder/QueryBuilderMethodDynamicReturnTypeExtension.php

@@ -2,28 +2,67 @@
 
 namespace PHPStan\Type\Doctrine\QueryBuilder;
 
+use PhpParser\Node;
 use PhpParser\Node\Expr\MethodCall;
 use PhpParser\Node\Identifier;
+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\MixedType;
 use PHPStan\Type\ObjectType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\TypeWithClassName;
 
-class QueryBuilderMethodDynamicReturnTypeExtension implements \PHPStan\Type\DynamicMethodReturnTypeExtension
+class QueryBuilderMethodDynamicReturnTypeExtension implements \PHPStan\Type\DynamicMethodReturnTypeExtension, BrokerAwareExtension
 {
 
 	private const MAX_COMBINATIONS = 16;
 
+	/** @var \PHPStan\DependencyInjection\Container */
+	private $container;
+
+	/** @var \PHPStan\Parser\Parser */
+	private $parser;
+
 	/** @var string|null */
 	private $queryBuilderClass;
 
-	public function __construct(?string $queryBuilderClass)
+	/** @var bool */
+	private $descendIntoOtherMethods;
+
+	/** @var \PHPStan\Broker\Broker */
+	private $broker;
+
+	public function __construct(
+		Container $container,
+		Parser $parser,
+		?string $queryBuilderClass,
+		bool $descendIntoOtherMethods
+	)
 	{
+		$this->container = $container;
+		$this->parser = $parser;
 		$this->queryBuilderClass = $queryBuilderClass;
+		$this->descendIntoOtherMethods = $descendIntoOtherMethods;
+	}
+
+	public function setBroker(Broker $broker): void
+	{
+		$this->broker = $broker;
 	}
 
 	public function getClass(): string
@@ -62,8 +101,16 @@ public function getTypeFromMethodCall(
 
 		$queryBuilderTypes = DoctrineTypeUtils::getQueryBuilderTypes($calledOnType);
 		if (count($queryBuilderTypes) === 0) {
-			return $calledOnType;
+			if (!$this->descendIntoOtherMethods || !$methodCall->var instanceof MethodCall) {
+				return $calledOnType;
+			}
+
+			$queryBuilderTypes = $this->findQueryBuilderTypesInCalledMethod($scope, $methodCall->var);
+			if (count($queryBuilderTypes) === 0) {
+				return $calledOnType;
+			}
 		}
+
 		if (count($queryBuilderTypes) > self::MAX_COMBINATIONS) {
 			return $calledOnType;
 		}
@@ -76,4 +123,137 @@ public function getTypeFromMethodCall(
 		return TypeCombinator::union(...$resultTypes);
 	}
 
+	/**
+	 * @param \PHPStan\Analyser\Scope $scope
+	 * @param \PhpParser\Node\Expr\MethodCall $methodCall
+	 * @return \PHPStan\Type\Doctrine\QueryBuilder\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 === false) {
+			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 \PHPStan\Analyser\NodeScopeResolver $nodeScopeResolver */
+		$nodeScopeResolver = $this->container->getByType(NodeScopeResolver::class);
+
+		/** @var \PHPStan\Analyser\ScopeFactory $scopeFactory */
+		$scopeFactory = $this->container->getByType(ScopeFactory::class);
+
+		$methodScope = $scopeFactory->create(
+			ScopeContext::create($fileName),
+			$scope->isDeclareStrictTypes(),
+			$methodReflection,
+			$scope->getNamespace()
+		)->enterClass($methodReflection->getDeclaringClass())->enterClassMethod($methodNode, [], null, null, false, false, false);
+
+		$queryBuilderTypes = [];
+
+		$nodeScopeResolver->processNodes($methodNode->stmts, $methodScope, 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 string $className
+	 * @param \PhpParser\Node[] $nodes
+	 * @return \PhpParser\Node\Stmt\Class_|null
+	 */
+	private function findClassNode(string $className, array $nodes): ?Class_
+	{
+		foreach ($nodes as $node) {
+			if (
+				$node instanceof Class_
+				&& $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 string $methodName
+	 * @param \PhpParser\Node\Stmt[] $classStatements
+	 * @return \PhpParser\Node\Stmt\ClassMethod|null
+	 */
+	private function findMethodNode(string $methodName, array $classStatements): ?ClassMethod
+	{
+		foreach ($classStatements as $statement) {
+			if (
+				$statement instanceof ClassMethod
+				&& $statement->name->toString() === $methodName
+			) {
+				return $statement;
+			}
+		}
+
+		return null;
+	}
+
 }

tests/Rules/Doctrine/ORM/QueryBuilderDqlRuleTest.php

@@ -4,6 +4,9 @@
 
 use Doctrine\ORM\Query\Expr\Base;
 use Doctrine\ORM\Query\Expr\OrderBy;
+use PHPStan\DependencyInjection\Container;
+use PHPStan\DependencyInjection\ContainerFactory;
+use PHPStan\DependencyInjection\Nette\NetteContainer;
 use PHPStan\Rules\Rule;
 use PHPStan\Testing\RuleTestCase;
 use PHPStan\Type\Doctrine\ArgumentsProcessor;
@@ -95,6 +98,10 @@ public function testRule(): void
 				'QueryBuilder: [Semantical Error] line 0, col 60 near \'transient = \': Error: Class PHPStan\Rules\Doctrine\ORM\MyEntity has no field or association named transient',
 				234,
 			],
+			[
+				'QueryBuilder: [Semantical Error] line 0, col 60 near \'nonexistent =\': Error: Class PHPStan\Rules\Doctrine\ORM\MyEntity has no field or association named nonexistent',
+				251,
+			],
 		]);
 	}
 
@@ -164,15 +171,26 @@ public function getDynamicMethodReturnTypeExtensions(): array
 	{
 		$objectMetadataResolver = new ObjectMetadataResolver(__DIR__ . '/entity-manager.php', null);
 		$argumentsProcessor = new ArgumentsProcessor();
+
 		return [
 			new CreateQueryBuilderDynamicReturnTypeExtension(null, $this->fasterVersion),
-			new QueryBuilderMethodDynamicReturnTypeExtension(null),
+			new QueryBuilderMethodDynamicReturnTypeExtension($this->getContainerWithDoctrineExtensions(), $this->getParser(), null, true),
 			new QueryBuilderGetQueryDynamicReturnTypeExtension($objectMetadataResolver, $argumentsProcessor, null),
 			new QueryGetDqlDynamicReturnTypeExtension(),
 			new ExpressionBuilderDynamicReturnTypeExtension($objectMetadataResolver, $argumentsProcessor),
 		];
 	}
 
+	private function getContainerWithDoctrineExtensions(): Container
+	{
+		$rootDir = __DIR__ . '/../../../../vendor/phpstan/phpstan';
+		$containerFactory = new ContainerFactory($rootDir);
+		return new NetteContainer($containerFactory->create($rootDir . '/tmp', [
+			$containerFactory->getConfigDirectory() . '/config.level7.neon',
+			__DIR__ . '/../../../../extension.neon',
+		], []));
+	}
+
 	/**
 	 * @return \PHPStan\Type\MethodTypeSpecifyingExtension[]
 	 */

tests/Rules/Doctrine/ORM/data/ClassWithQueryBuilder.php

@@ -0,0 +1,28 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Rules\Doctrine\ORM;
+
+use Doctrine\ORM\EntityManager;
+use Doctrine\ORM\QueryBuilder;
+
+class ClassWithQueryBuilder
+{
+
+	/** @var EntityManager */
+	private $entityManager;
+
+	public function __construct(EntityManager $entityManager)
+	{
+		$this->entityManager = $entityManager;
+	}
+
+	public function getQueryBuilder(): QueryBuilder
+	{
+		$queryBuilder = $this->entityManager->createQueryBuilder();
+		$queryBuilder->select('e')
+			->from(MyEntity::class, 'e');
+
+		return $queryBuilder;
+	}
+
+}

tests/Rules/Doctrine/ORM/data/query-builder-dql.php

@@ -244,4 +244,11 @@ public function weirdTypeSpecifyingExtensionProblemCorrect(): void
 		$queryBuilder->getQuery();
 	}
 
+	public function queryBuilderFromSomewhereElse(): void
+	{
+		$class = new ClassWithQueryBuilder($this->entityManager);
+		$queryBuilder = $class->getQueryBuilder()->andWhere('e.nonexistent = :test');
+		$queryBuilder->getQuery();
+	}
+
 }