Dynamic return type extension for EntityRepository methods via a generic type created by EM::getRepository()

Author: ondrejmirtes

extension.neon

@@ -11,3 +11,11 @@ services:
 		class: PHPStan\Type\Doctrine\EntityManagerFindDynamicReturnTypeExtension
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: PHPStan\Type\Doctrine\EntityManagerGetRepositoryDynamicReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: PHPStan\Type\Doctrine\EntityRepositoryDynamicReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension

src/Type/Doctrine/EntityManagerGetRepositoryDynamicReturnTypeExtension.php

@@ -0,0 +1,54 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine;
+
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Type\Type;
+
+class EntityManagerGetRepositoryDynamicReturnTypeExtension implements \PHPStan\Type\DynamicMethodReturnTypeExtension
+{
+
+	public static function getClass(): string
+	{
+		return \Doctrine\ORM\EntityManager::class;
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		return $methodReflection->getName() === 'getRepository';
+	}
+
+	public function getTypeFromMethodCall(
+		MethodReflection $methodReflection,
+		MethodCall $methodCall,
+		Scope $scope
+	): Type
+	{
+		if (count($methodCall->args) === 0) {
+			return $methodReflection->getReturnType();
+		}
+		$arg = $methodCall->args[0]->value;
+		if (!($arg instanceof \PhpParser\Node\Expr\ClassConstFetch)) {
+			return $methodReflection->getReturnType();
+		}
+
+		$class = $arg->class;
+		if (!($class instanceof \PhpParser\Node\Name)) {
+			return $methodReflection->getReturnType();
+		}
+
+		$class = (string) $class;
+		if ($class === 'static') {
+			return $methodReflection->getReturnType();
+		}
+
+		if ($class === 'self') {
+			$class = $scope->getClassReflection()->getName();
+		}
+
+		return new EntityRepositoryType($class);
+	}
+
+}

src/Type/Doctrine/EntityRepositoryDynamicReturnTypeExtension.php

@@ -0,0 +1,51 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine;
+
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Type\ArrayType;
+use PHPStan\Type\MixedType;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+
+class EntityRepositoryDynamicReturnTypeExtension implements \PHPStan\Type\DynamicMethodReturnTypeExtension
+{
+
+	public static function getClass(): string
+	{
+		return \Doctrine\ORM\EntityRepository::class;
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		$methodName = $methodReflection->getName();
+		return strpos($methodName, 'findBy') === 0
+			|| strpos($methodName, 'findOneBy') === 0
+			|| $methodName === 'findAll'
+			|| $methodName === 'find';
+	}
+
+	public function getTypeFromMethodCall(
+		MethodReflection $methodReflection,
+		MethodCall $methodCall,
+		Scope $scope
+	): Type
+	{
+		$calledOnType = $scope->getType($methodCall->var);
+		if (!$calledOnType instanceof EntityRepositoryType) {
+			return new MixedType();
+		}
+		$methodName = $methodReflection->getName();
+		$entityType = new ObjectType($calledOnType->getEntityClass());
+
+		if ($methodName === 'find' || strpos($methodName, 'findOneBy') === 0) {
+			return TypeCombinator::addNull($entityType);
+		}
+
+		return new ArrayType($entityType);
+	}
+
+}

src/Type/Doctrine/EntityRepositoryType.php

@@ -0,0 +1,31 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine;
+
+use PHPStan\Type\ObjectType;
+
+class EntityRepositoryType extends ObjectType
+{
+
+	/**
+	 * @var string
+	 */
+	private $entityClass;
+
+	public function __construct(string $entityClass)
+	{
+		parent::__construct(\Doctrine\ORM\EntityRepository::class);
+		$this->entityClass = $entityClass;
+	}
+
+	public function getEntityClass(): string
+	{
+		return $this->entityClass;
+	}
+
+	public function describe(): string
+	{
+		return sprintf('%s<%s>', parent::describe(), $this->entityClass);
+	}
+
+}