From f42f93f75ada7f577099e0e157faa8f3f3a0721c Mon Sep 17 00:00:00 2001
From: Maciej Malarz <malarzm@gmail.com>
Date: Wed, 13 Dec 2017 13:24:27 +0100
Subject: [PATCH] Resolve custom repository classes

---
 extension.neon                                |  9 ++-
 src/DoctrineClassMetadataProvider.php         | 81 +++++++++++++++++++
 ...etRepositoryDynamicReturnTypeExtension.php | 25 ++++--
 3 files changed, 106 insertions(+), 9 deletions(-)
 create mode 100644 src/DoctrineClassMetadataProvider.php

diff --git a/extension.neon b/extension.neon
index 21e8e248..a83f504f 100644
--- a/extension.neon
+++ b/extension.neon
@@ -1,6 +1,7 @@
 parameters:
 	doctrine:
 		repositoryClass: Doctrine\ORM\EntityRepository
+		metadata: []
 
 services:
 	-
@@ -18,7 +19,7 @@ services:
 	-
 		class: PHPStan\Type\Doctrine\EntityManagerGetRepositoryDynamicReturnTypeExtension
 		arguments:
-			repositoryClass: %doctrine.repositoryClass%
+			metadataProvider: @PHPStan\DoctrineClassMetadataProvider
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
 	-
@@ -29,3 +30,9 @@ services:
 		class: PHPStan\Type\Doctrine\EntityRepositoryDynamicReturnTypeExtension
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: PHPStan\DoctrineClassMetadataProvider
+		arguments:
+			repositoryClass: %doctrine.repositoryClass%
+			mapping: %doctrine.metadata%
+
diff --git a/src/DoctrineClassMetadataProvider.php b/src/DoctrineClassMetadataProvider.php
new file mode 100644
index 00000000..41169ea7
--- /dev/null
+++ b/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);
+    }
+    
+}
diff --git a/src/Type/Doctrine/EntityManagerGetRepositoryDynamicReturnTypeExtension.php b/src/Type/Doctrine/EntityManagerGetRepositoryDynamicReturnTypeExtension.php
index 3be6d943..87a7c0c9 100644
--- a/src/Type/Doctrine/EntityManagerGetRepositoryDynamicReturnTypeExtension.php
+++ b/src/Type/Doctrine/EntityManagerGetRepositoryDynamicReturnTypeExtension.php
@@ -2,23 +2,27 @@
 
 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\Type\Type;
 
 class EntityManagerGetRepositoryDynamicReturnTypeExtension implements \PHPStan\Type\DynamicMethodReturnTypeExtension
 {
 
-	/**
-	 * @var string
-	 */
-	private $repositoryClass;
+    /**
+     * @var DoctrineClassMetadataProvider
+     */
+    private $metadataProvider;
 
-	public function __construct(string $repositoryClass)
+    public function __construct(DoctrineClassMetadataProvider $metadataProvider)
 	{
-		$this->repositoryClass = $repositoryClass;
-	}
+        $this->metadataProvider = $metadataProvider;
+        AnnotationRegistry::registerLoader('class_exists');
+    }
 
 	public function getClass(): string
 	{
@@ -58,7 +62,12 @@ public function getTypeFromMethodCall(
 			$class = $scope->getClassReflection()->getName();
 		}
 
-		return new EntityRepositoryType($class, $this->repositoryClass);
+        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());
+        }
 	}
 
 }