Add support for union of class string

Author: VincentLanglet

src/Type/Doctrine/GetRepositoryDynamicReturnTypeExtension.php

@@ -21,6 +21,7 @@
 use PHPStan\Type\ObjectType;
 use PHPStan\Type\ObjectWithoutClassType;
 use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
 use function count;
 
 class GetRepositoryDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
@@ -96,22 +97,26 @@ public function getTypeFromMethodCall(
 
 		$classType = $argType->getClassStringObjectType();
 		$objectNames = $classType->getObjectClassNames();
-		if (count($objectNames) !== 1) {
+
+		if (count($objectNames) === 0) {
 			return new GenericObjectType(
 				$defaultRepositoryClass,
 				[$classType]
 			);
 		}
 
-		try {
-			$repositoryClass = $this->getRepositoryClass($objectNames[0], $defaultRepositoryClass);
-		} catch (\Doctrine\Persistence\Mapping\MappingException | MappingException | AnnotationException $e) {
-			return $this->getDefaultReturnType($scope, $methodCall->getArgs(), $methodReflection, $defaultRepositoryClass);
+		$repositoryTypes = [];
+		foreach ($objectNames as $objectName) {
+			try {
+				$repositoryClass = $this->getRepositoryClass($objectName, $defaultRepositoryClass);
+			} catch (\Doctrine\Persistence\Mapping\MappingException | MappingException | AnnotationException $e) {
+				return $this->getDefaultReturnType($scope, $methodCall->getArgs(), $methodReflection, $defaultRepositoryClass);
+			}
+
+			$repositoryTypes[] = new GenericObjectType($repositoryClass, [$classType]);
 		}
 
-		return new GenericObjectType($repositoryClass, [
-			$classType,
-		]);
+		return TypeCombinator::union(...$repositoryTypes);
 	}
 
 	/**

tests/Rules/Doctrine/ORM/MagicRepositoryMethodCallRuleTest.php

@@ -88,6 +88,14 @@ public function testRule(): void
 				'Method Doctrine\Persistence\ObjectRepository<PHPStan\Rules\Doctrine\ORM\MyEntity>::countByTitle() invoked with 2 parameters, 1 required.',
 				85,
 			],
+			[
+				'Call to an undefined method PHPStan\Rules\Doctrine\ORM\TestRepository<PHPStan\Rules\Doctrine\ORM\MySecondEntity|PHPStan\Rules\Doctrine\ORM\MyThirdEntity>::findByTransient().',
+				97,
+			],
+			[
+				'Call to an undefined method PHPStan\Rules\Doctrine\ORM\TestRepository<PHPStan\Rules\Doctrine\ORM\MySecondEntity|PHPStan\Rules\Doctrine\ORM\MyThirdEntity>::findByNonexistent().',
+				98,
+			],
 		]);
 	}
 

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

@@ -0,0 +1,32 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Rules\Doctrine\ORM;
+
+use Doctrine\ORM\Mapping as ORM;
+
+/**
+ * @ORM\Entity(repositoryClass="PHPStan\Rules\Doctrine\ORM\TestRepository")
+ */
+class MyThirdEntity
+{
+	/**
+	 * @ORM\Id()
+	 * @ORM\GeneratedValue()
+	 * @ORM\Column(type="integer")
+	 *
+	 * @var int
+	 */
+	private $id;
+
+	/**
+	 * @var string
+	 * @ORM\Column(type="string")
+	 */
+	private $title;
+
+	/**
+	 * @var string
+	 */
+	private $transient;
+
+}

tests/Rules/Doctrine/ORM/data/magic-repository.php

@@ -84,4 +84,18 @@ public function doCountByWithOptionals(): void
 		$entityRepository = $this->entityManager->getRepository(MyEntity::class);
 		$entityRepository->countByTitle('test', ['id' => 'DESC']);
 	}
+
+	/**
+	 * @param class-string<MySecondEntity|MyThirdEntity> $entityClass
+	 */
+	public function doFindByWithRepositoryOfMultiplesClasses(string $entityClass): void
+	{
+		$entityRepository = $this->entityManager->getRepository($entityClass);
+		$entityRepository->findBy(['id' => 1]);
+		$entityRepository->findById(1);
+		$entityRepository->findByTitle('test');
+		$entityRepository->findByTransient('test');
+		$entityRepository->findByNonexistent('test');
+		$entityRepository->findByCustomMethod();
+	}
 }