Resolve custom repository classes

Author: malarzm

extension.neon

@@ -1,6 +1,7 @@
 parameters:
 	doctrine:
 		repositoryClass: Doctrine\ORM\EntityRepository
+		metadata: []
 		allCollectionsSelectable: true
 
 conditionalTags:
@@ -21,7 +22,7 @@ services:
 	-
 		class: PHPStan\Type\Doctrine\ObjectManagerGetRepositoryDynamicReturnTypeExtension
 		arguments:
-			repositoryClass: %doctrine.repositoryClass%
+			metadataProvider: @PHPStan\DoctrineClassMetadataProvider
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
 	-
@@ -32,3 +33,9 @@ services:
 		class: PHPStan\Type\Doctrine\ObjectRepositoryDynamicReturnTypeExtension
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: PHPStan\DoctrineClassMetadataProvider
+		arguments:
+			repositoryClass: %doctrine.repositoryClass%
+			mapping: %doctrine.metadata%
+

src/DoctrineClassMetadataProvider.php

@@ -0,0 +1,81 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan;
+
+use Doctrine\Common\Annotations\AnnotationReader;
+use Doctrine\Common\EventManager;
+use Doctrine\Common\Persistence\Mapping\Driver\MappingDriver;
+use Doctrine\Common\Persistence\Mapping\Driver\MappingDriverChain;
+use Doctrine\ORM\Configuration;
+use Doctrine\ORM\EntityManager;
+use Doctrine\ORM\Mapping;
+use Doctrine\ORM\Proxy\ProxyFactory;
+
+class DoctrineClassMetadataProvider
+{
+
+    /**
+     * @var EntityManager
+     */
+    private $em;
+
+    /**
+     * @var string
+     */
+    private $repositoryClass;
+
+    public function __construct(string $repositoryClass, array $mapping)
+    {
+        $configuration = new Configuration();
+        $configuration->setDefaultRepositoryClassName($repositoryClass);
+        $configuration->setMetadataDriverImpl($this->setupMappingDriver($mapping));
+        $configuration->setProxyDir('/dev/null');
+        $configuration->setProxyNamespace('__DP__');
+        $configuration->setAutoGenerateProxyClasses(ProxyFactory::AUTOGENERATE_EVAL);
+        $evm = new EventManager();
+        $this->em = EntityManager::create(
+            \Doctrine\DBAL\DriverManager::getConnection(['host' => '/:memory:', 'driver' => 'pdo_sqlite'], $configuration, $evm),
+            $configuration,
+            $evm
+        );
+        $this->repositoryClass = $repositoryClass;
+    }
+
+    private function setupMappingDriver(array $mapping): MappingDriver
+    {
+        $driver = new MappingDriverChain();
+        foreach ($mapping as $namespace => $config) {
+            switch ($config['type']) {
+                case 'annotation':
+                    $nested = new Mapping\Driver\AnnotationDriver(new AnnotationReader(), $config['paths']);
+                    break;
+                case 'yml':
+                case 'yaml':
+                    $nested = new Mapping\Driver\YamlDriver($config['paths']);
+                    break;
+                case 'xml':
+                    $nested = new Mapping\Driver\XmlDriver($config['paths']);
+                    break;
+                default:
+                    throw new \InvalidArgumentException('Unknown mapping type: ' . $config['type']);
+            }
+            $driver->addDriver($nested, $namespace);
+        }
+        return $driver;
+    }
+
+    public function getBaseRepositoryClass(): string
+    {
+        return $this->repositoryClass;
+    }
+
+    /**
+     * @throws \Doctrine\Common\Persistence\Mapping\MappingException
+     * @throws \ReflectionException
+     */
+    public function getMetadataFor(string $className): Mapping\ClassMetadataInfo
+    {
+        return $this->em->getClassMetadata($className);
+    }
+    
+}

src/Type/Doctrine/ObjectManagerGetRepositoryDynamicReturnTypeExtension.php

@@ -2,8 +2,11 @@
 
 namespace PHPStan\Type\Doctrine;
 
+use Doctrine\Common\Annotations\AnnotationRegistry;
+use Doctrine\Common\Persistence\Mapping\MappingException;
 use PhpParser\Node\Expr\MethodCall;
 use PHPStan\Analyser\Scope;
+use PHPStan\DoctrineClassMetadataProvider;
 use PHPStan\Reflection\MethodReflection;
 use PHPStan\Reflection\ParametersAcceptorSelector;
 use PHPStan\Type\Constant\ConstantStringType;
@@ -16,10 +19,16 @@ class ObjectManagerGetRepositoryDynamicReturnTypeExtension implements \PHPStan\T
 	/** @var string */
 	private $repositoryClass;
 
-	public function __construct(string $repositoryClass)
+    /**
+     * @var DoctrineClassMetadataProvider
+     */
+    private $metadataProvider;
+
+    public function __construct(DoctrineClassMetadataProvider $metadataProvider)
 	{
-		$this->repositoryClass = $repositoryClass;
-	}
+        $this->metadataProvider = $metadataProvider;
+        AnnotationRegistry::registerLoader('class_exists');
+    }
 
 	public function getClass(): string
 	{
@@ -48,6 +57,27 @@ public function getTypeFromMethodCall(
 		}
 
 		return new ObjectRepositoryType($argType->getValue(), $this->repositoryClass);
+
+		$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();
+		}
+
+        try {
+            $metadata = $this->metadataProvider->getMetadataFor($class);
+            return new EntityRepositoryType($class, $metadata->customRepositoryClassName ?: $this->metadataProvider->getBaseRepositoryClass());
+        } catch (MappingException $e) {
+		    return new EntityRepositoryType($class, $this->metadataProvider->getBaseRepositoryClass());
+        }
 	}
 
 }