Allow custom getRepository()-like methods

src/main/java/fr/adrienbrault/idea/symfony2plugin/doctrine/ObjectRepositoryTypeProvider.java

@@ -4,10 +4,7 @@
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
 import com.jetbrains.php.PhpIndex;
-import com.jetbrains.php.lang.psi.elements.Method;
-import com.jetbrains.php.lang.psi.elements.MethodReference;
-import com.jetbrains.php.lang.psi.elements.PhpClass;
-import com.jetbrains.php.lang.psi.elements.PhpNamedElement;
+import com.jetbrains.php.lang.psi.elements.*;
 import com.jetbrains.php.lang.psi.resolve.types.PhpType;
 import com.jetbrains.php.lang.psi.resolve.types.PhpTypeProvider3;
 import fr.adrienbrault.idea.symfony2plugin.Settings;
@@ -29,6 +26,8 @@ public class ObjectRepositoryTypeProvider implements PhpTypeProvider3 {
         new MethodMatcher.CallToSignature("\\Doctrine\\Common\\Persistence\\ObjectManager", "getRepository"),
     };
 
+    private static String repositoryClass = "\\Doctrine\\Common\\Persistence\\ObjectRepository";
+
     final public static char TRIM_KEY = '\u0185';
 
     @Override
@@ -43,11 +42,10 @@ public PhpType getType(PsiElement e) {
             return null;
         }
 
-        if(!(e instanceof MethodReference) || !PhpElementsUtil.isMethodWithFirstStringOrFieldReference(e, "getRepository")) {
+        if(!(e instanceof MethodReference)) {
             return null;
         }
 
-
         String refSignature = ((MethodReference)e).getSignature();
         if(StringUtil.isEmpty(refSignature)) {
             return null;
@@ -80,7 +78,13 @@ public Collection<? extends PhpNamedElement> getBySignature(String expression, S
             return phpNamedElementCollections;
         }
 
-        if (!PhpElementsUtil.isMethodInstanceOf((Method) phpNamedElement, GET_REPOSITORIES_SIGNATURES)) {
+        PhpReturnType returnType = ((Method) phpNamedElement).getReturnType();
+        String returnTypeName = returnType == null ? null : returnType.getType().toString();
+
+        boolean isMethod = PhpElementsUtil.isMethodInstanceOf((Method) phpNamedElement, GET_REPOSITORIES_SIGNATURES);
+        boolean isReturnType = returnType != null && PhpElementsUtil.isInstanceOf(project, returnTypeName, repositoryClass);
+
+        if (!isMethod && !isReturnType) {
             return phpNamedElementCollections;
         }
 

src/test/java/fr/adrienbrault/idea/symfony2plugin/tests/doctrine/ObjectRepositoryTypeProviderTest.java

@@ -45,7 +45,6 @@ public void testGetRepositoryResolveByRepository() {
                 "$em->getRepository(\\Foo\\Bar::class)->b<caret>ar();",
             PlatformPatterns.psiElement(Method.class).withName("bar")
         );
-
     }
 
     /**
@@ -61,6 +60,55 @@ public void testGetRepositoryResolveByRepositoryApiClassConstantCompatibility()
         );
     }
 
+    /**
+     * @see fr.adrienbrault.idea.symfony2plugin.doctrine.ObjectRepositoryTypeProvider
+     */
+    public void testGetRepositoryResolveByReturnTypeMethod() {
+        assertPhpReferenceResolveTo(PhpFileType.INSTANCE,
+                "<?php" +
+                        "/** @var \\Foo\\BarController $cont */\n" +
+                        "$cont->getRepo('\\Foo\\Bar')->b<caret>ar();",
+                PlatformPatterns.psiElement(Method.class).withName("bar")
+        );
+
+        assertPhpReferenceSignatureContains(PhpFileType.INSTANCE,
+                "<?php" +
+                        "/** @var \\Foo\\BarController $cont */\n" +
+                        "$cont->getRepo('\\Foo\\Bar')->b<caret>ar();",
+                "#M#" + '\u0151' + "#M#C\\Foo\\BarController.getRepo" + '\u0185' + "\\Foo\\Bar.bar"
+        );
+
+        assertPhpReferenceResolveTo(PhpFileType.INSTANCE,
+                "<?php" +
+                        "/** @var \\Foo\\BarController $cont */\n" +
+                        "$cont->getRepo(\\Foo\\Bar::class)->b<caret>ar();",
+                PlatformPatterns.psiElement(Method.class).withName("bar")
+        );
+
+    }
+
+    public void testGetRepositoryNotResolveByWrongMethod() {
+        assertPhpReferenceNotResolveTo(PhpFileType.INSTANCE,
+                "<?php" +
+                        "/** @var \\Foo\\BarController $cont */\n" +
+                        "$cont->getBoo(\\Foo\\Bar::class)->b<caret>ar();",
+                PlatformPatterns.psiElement(Method.class).withName("bar")
+        );
+    }
+
+    /**
+     * @see fr.adrienbrault.idea.symfony2plugin.doctrine.ObjectRepositoryTypeProvider
+     */
+    public void testGetRepositoryResolveByReturnTypeMethodApiClassConstantCompatibility() {
+        String result = "#M#" + '\u0151' + "#M#C\\Foo\\BarController.getRepo" + '\u0185' + "#K#C\\Foo\\Bar.class.bar";
+
+        assertPhpReferenceSignatureContains(PhpFileType.INSTANCE, "<?php" +
+                        "/** @var \\Foo\\BarController $cont */\n" +
+                        "$cont->getRepo(\\Foo\\Bar::class)->b<caret>ar();",
+                result
+        );
+    }
+
     /**
      * @see fr.adrienbrault.idea.symfony2plugin.doctrine.ObjectRepositoryTypeProvider
      */

src/test/java/fr/adrienbrault/idea/symfony2plugin/tests/doctrine/fixtures/classes.php

@@ -11,14 +11,20 @@ function getRepository();
 }
 
 namespace Foo {
-
     use Doctrine\Common\Persistence\ObjectRepository;
 
     abstract class Bar implements ObjectRepository {}
 
-    class BarRepository {
+    class BarRepository implements ObjectRepository {
         public function bar() {}
     }
+
+    class WrongClass {}
+
+    class BarController {
+        public function getRepo($name): BarRepository {}
+        public function getBoo($name): WrongClass {}
+    }
 }
 
 namespace Doctrine\Common\Persistence