support iterator in twig loop completion #1097 #1035

Author: Haehnchen

src/fr/adrienbrault/idea/symfony2plugin/templating/util/TwigTypeResolveUtil.java

@@ -37,6 +37,7 @@
 import java.util.*;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 /**
  * @author Daniel Espendiller <daniel@espendiller.net>
@@ -259,7 +260,13 @@ public static Map<String, PsiVariable> collectScopeVariables(@NotNull PsiElement
         globalVars.putAll(convertHashMapToTypeSet(findInlineStatementVariableDocBlock(psiElement, TwigElementTypes.FOR_STATEMENT)));
 
         for(Map.Entry<String, Set<String>> entry: globalVars.entrySet()) {
-            controllerVars.put(entry.getKey(), new PsiVariable(entry.getValue(), null));
+            Set<String> types = entry.getValue();
+
+            // collect iterator
+            types.addAll(collectIteratorReturns(psiElement, entry.getValue()));
+
+            // convert to variable model
+            controllerVars.put(entry.getKey(), new PsiVariable(types, null));
         }
 
         // check if we are in "for" scope and resolve types ending with []
@@ -268,6 +275,47 @@ public static Map<String, PsiVariable> collectScopeVariables(@NotNull PsiElement
         return controllerVars;
     }
 
+    /**
+     * Extract magic iterator implementation like "getIterator" or "__iterator"
+     *
+     * "@TODO find core stuff for resolve possible class return values"
+     *
+     * "getIterator", "@method Foo __iterator", "@method Foo[] __iterator"
+     */
+    @NotNull
+    private static Set<String> collectIteratorReturns(@NotNull PsiElement psiElement, @NotNull Set<String> types) {
+        Set<String> arrayValues = new HashSet<>();
+        for (String type : types) {
+            PhpClass phpClass = PhpElementsUtil.getClassInterface(psiElement.getProject(), type);
+
+            if(phpClass == null) {
+                continue;
+            }
+
+            for (String methodName : new String[]{"getIterator", "__iterator", "current"}) {
+                Method method = phpClass.findMethodByName(methodName);
+                if(method != null) {
+                    // @method Foo __iterator
+                    // @method Foo[] __iterator
+                    Set<String> iteratorTypes = method.getType().getTypes();
+                    if("__iterator".equals(methodName) || "current".equals(methodName)) {
+                        arrayValues.addAll(iteratorTypes.stream().map(x ->
+                            !x.endsWith("[]") ? x + "[]" : x
+                        ).collect(Collectors.toSet()));
+                    } else {
+                        // Foobar[]
+                        for (String iteratorType : iteratorTypes) {
+                            if(iteratorType.endsWith("[]")) {
+                                arrayValues.add(iteratorType);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        return arrayValues;
+    }
 
     private static Collection<String> collectForArrayScopeVariablesFoo(Project project, String[] typeName, PsiVariable psiVariable) {
 

tests/fr/adrienbrault/idea/symfony2plugin/tests/templating/variable/collector/FileDocVariableCollectorTest.java

@@ -1,5 +1,6 @@
 package fr.adrienbrault.idea.symfony2plugin.tests.templating.variable.collector;
 
+import com.jetbrains.php.lang.PhpFileType;
 import com.jetbrains.twig.TwigFileType;
 import fr.adrienbrault.idea.symfony2plugin.tests.SymfonyLightCodeInsightFixtureTestCase;
 
@@ -25,9 +26,10 @@ public void setUp() throws Exception {
             "  private function privateBar() {}\n" +
             "  /** @return FooClass[] */\n" +
             "  public function getNested() {}\n" +
+            "  /** @return FooClass[] */\n" +
+            "  public function getIterator() {}\n" +
             "}"
         );
-
     }
 
     public void testFileBasedVarDocPhpTypes() {
@@ -91,7 +93,70 @@ public void testVarArrayIteration() {
                 "{% endfor %}\n"
             , "fooBar"
         );
+    }
+
+    public void testVarArrayIterationViaIterationClassImplementations() {
+        myFixture.configureByText("class1.php", "<?php\n" +
+            "namespace Bar;\n" +
+            "/**\n" +
+            " * @method FooClassIteratorArray[] __iterator\n" +
+            " */\n" +
+            "class FooClassIteratorArray {\n" +
+            "  public function getFooBar() {}\n" +
+            "}"
+        );
+
+        myFixture.configureByText("class2.php", "<?php\n" +
+            "namespace Bar;\n" +
+            "/**\n" +
+            " * @method FooClassIterator __iterator\n" +
+            " */\n" +
+            "class FooClassIterator {\n" +
+            "  public function getFooBar() {}\n" +
+            "}"
+        );
+
+        myFixture.configureByText("class3.php", "<?php\n" +
+            "namespace Bar;\n" +
+            "/**\n" +
+            " * @method FooClassCurrent current\n" +
+            " */\n" +
+            "class FooClassCurrent {\n" +
+            "  public function getFooBar() {}\n" +
+            "}"
+        );
+
+        assertCompletionContains(TwigFileType.INSTANCE, "" +
+                "{# @var bars \\Bar\\FooClass #}\n" +
+                "{% for bar in bars %}\n" +
+                "  {{ bar.<caret> }}\n" +
+                "{% endfor %}\n"
+            , "fooBar"
+        );
 
+        assertCompletionContains(TwigFileType.INSTANCE, "" +
+                "{# @var bars \\Bar\\FooClassIteratorArray #}\n" +
+                "{% for bar in bars %}\n" +
+                "  {{ bar.<caret> }}\n" +
+                "{% endfor %}\n"
+            , "fooBar"
+        );
+
+        assertCompletionContains(TwigFileType.INSTANCE, "" +
+                "{# @var bars \\Bar\\FooClassIterator #}\n" +
+                "{% for bar in bars %}\n" +
+                "  {{ bar.<caret> }}\n" +
+                "{% endfor %}\n"
+            , "fooBar"
+        );
+
+        assertCompletionContains(TwigFileType.INSTANCE, "" +
+                "{# @var bars \\Bar\\FooClassCurrent #}\n" +
+                "{% for bar in bars %}\n" +
+                "  {{ bar.<caret> }}\n" +
+                "{% endfor %}\n"
+            , "fooBar"
+        );
     }
 
     /**