OtherMethodQueryBuilderParser: fix for BranchingQueryBuilderType

Author: janedbal

src/Type/Doctrine/QueryBuilder/OtherMethodQueryBuilderParser.php

@@ -19,6 +19,10 @@
 use PHPStan\Parser\Parser;
 use PHPStan\Reflection\ReflectionProvider;
 use PHPStan\Type\Generic\TemplateTypeMap;
+use PHPStan\Type\IntersectionType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeTraverser;
+use PHPStan\Type\UnionType;
 use function count;
 use function is_array;
 
@@ -118,11 +122,18 @@ private function findQueryBuilderTypesInCalledMethod(Scope $scope, MethodCall $m
 			}
 
 			$exprType = $scope->getType($node->expr);
-			if (!$exprType instanceof QueryBuilderType) {
-				return;
-			}
 
-			$queryBuilderTypes[] = $exprType;
+			TypeTraverser::map($exprType, static function (Type $type, callable $traverse) use (&$queryBuilderTypes): Type {
+				if ($type instanceof UnionType || $type instanceof IntersectionType) {
+					return $traverse($type);
+				}
+
+				if ($type instanceof QueryBuilderType) {
+					$queryBuilderTypes[] = $type;
+				}
+
+				return $type;
+			});
 		});
 
 		return $queryBuilderTypes;

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

@@ -140,10 +140,11 @@ public function testQueryResultTypeIsVoidWithDeleteOrUpdate(EntityManagerInterfa
 
 	public function testQueryTypeIsInferredOnAcrossMethods(EntityManagerInterface $em): void
 	{
-		$query = $this->getQueryBuilder($em)
-			->getQuery();
+		$query = $this->getQueryBuilder($em)->getQuery();
+		$branchingQuery = $this->getBranchingQueryBuilder($em)->getQuery();
 
 		assertType('Doctrine\ORM\Query<null, QueryResult\Entities\Many>', $query);
+		assertType('Doctrine\ORM\Query<null, QueryResult\Entities\Many>', $branchingQuery);
 	}
 
 	private function getQueryBuilder(EntityManagerInterface $em): QueryBuilder
@@ -152,4 +153,17 @@ private function getQueryBuilder(EntityManagerInterface $em): QueryBuilder
 			->select('m')
 			->from(Many::class, 'm');
 	}
+
+	private function getBranchingQueryBuilder(EntityManagerInterface $em): QueryBuilder
+	{
+		$queryBuilder = $em->createQueryBuilder()
+			->select('m')
+			->from(Many::class, 'm');
+
+		if (random_int(0, 1) === 1) {
+			$queryBuilder->andWhere('m.intColumn = 1');
+		}
+
+		return $queryBuilder;
+	}
 }