[python] postfix templates: add .fore/.itere/.enum (PY-27374)

GitOrigin-RevId: 1e389ba4ef18ebb13da8665fb5a47708cddc6ccf

python/src/com/jetbrains/python/codeInsight/postfix/PyForEnumeratePostfixTemplate.kt

@@ -0,0 +1,23 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.codeInsight.postfix
+
+import com.intellij.codeInsight.template.Template
+import com.intellij.codeInsight.template.impl.MacroCallNode
+import com.intellij.codeInsight.template.impl.TextExpression
+import com.intellij.codeInsight.template.impl.VariableNode
+import com.intellij.codeInsight.template.postfix.templates.PostfixTemplateProvider
+import com.intellij.psi.PsiElement
+import com.jetbrains.python.codeInsight.liveTemplates.CollectionElementNameMacro
+
+class PyForEnumeratePostfixTemplate(name: String, provider: PostfixTemplateProvider) : PyEditablePostfixTemplate(
+  name, name, "for \$INDEX$, \$VAR$ in enumerate(\$EXPR$):\n    \$END$", "for i, e in enumerate(expr)",
+  setOf(PyPostfixTemplateExpressionCondition.PyIterable()), false, provider, true) {
+
+  override fun addTemplateVariables(element: PsiElement, template: Template) {
+    super.addTemplateVariables(element, template)
+    val name = MacroCallNode(CollectionElementNameMacro())
+    name.addParameter(VariableNode("EXPR", null))
+    template.addVariable("VAR", name, TextExpression("e"), true)
+    template.addVariable("INDEX", null, TextExpression("i"), true)
+  }
+}

python/src/com/jetbrains/python/codeInsight/postfix/PyPostfixTemplateProvider.java

@@ -29,6 +29,9 @@ public class PyPostfixTemplateProvider implements PostfixTemplateProvider {
     new PyWhilePostfixTemplate(this),
     new PyForPostfixTemplate("for", this),
     new PyForPostfixTemplate("iter", this),
+    new PyForEnumeratePostfixTemplate("fore", this),
+    new PyForEnumeratePostfixTemplate("itere", this),
+    new PyForEnumeratePostfixTemplate("enum", this),
     new PyIsNonePostfixTemplate(this),
     new PyIsNotNonePostfixTemplate(this),
     new PyIsNotNonePostfixTemplate(this),

python/testData/postfix/enum/complexExpression.py

@@ -0,0 +1,2 @@
+def f(a:list, b:list, c:int):
+    (2*(a + b)) * c.enum<caret>

python/testData/postfix/enum/complexExpression_after.py

@@ -0,0 +1,3 @@
+def f(a:list, b:list, c:int):
+    for i, e in enumerate((2 * (a + b)) * c):
+        <caret>

python/testData/postfix/enum/function.py

@@ -0,0 +1,2 @@
+def f(a:list):
+    a.enum<caret>

python/testData/postfix/enum/function_after.py

@@ -0,0 +1,3 @@
+def f(a:list):
+    for i, e in enumerate(a):
+        <caret>

python/testData/postfix/enum/topLevel.py

@@ -0,0 +1 @@
+[].enum<caret>

python/testData/postfix/enum/topLevel_after.py

@@ -0,0 +1,2 @@
+for i, e in enumerate([]):
+    <caret>

python/testData/postfix/fore/complexExpression.py

@@ -0,0 +1,2 @@
+def f(a:list, b:list, c:int):
+    (2*(a + b)) * c.fore<caret>

python/testData/postfix/fore/complexExpression_after.py

@@ -0,0 +1,3 @@
+def f(a:list, b:list, c:int):
+    for i, e in enumerate((2 * (a + b)) * c):
+        <caret>

python/testData/postfix/fore/function.py

@@ -0,0 +1,2 @@
+def f(a:list):
+    a.fore<caret>

python/testData/postfix/fore/function_after.py

@@ -0,0 +1,3 @@
+def f(a:list):
+    for i, e in enumerate(a):
+        <caret>

python/testData/postfix/fore/topLevel.py

@@ -0,0 +1 @@
+[].fore<caret>

python/testData/postfix/fore/topLevel_after.py

@@ -0,0 +1,2 @@
+for i, e in enumerate([]):
+    <caret>

python/testData/postfix/itere/complexExpression.py

@@ -0,0 +1,2 @@
+def f(a:list, b:list, c:int):
+    (2*(a + b)) * c.itere<caret>

python/testData/postfix/itere/complexExpression_after.py

@@ -0,0 +1,3 @@
+def f(a:list, b:list, c:int):
+    for i, e in enumerate((2 * (a + b)) * c):
+        <caret>

python/testData/postfix/itere/function.py

@@ -0,0 +1,2 @@
+def f(a:list):
+    a.itere<caret>

python/testData/postfix/itere/function_after.py

@@ -0,0 +1,3 @@
+def f(a:list):
+    for i, e in enumerate(a):
+        <caret>

python/testData/postfix/itere/topLevel.py

@@ -0,0 +1 @@
+[].itere<caret>

python/testData/postfix/itere/topLevel_after.py

@@ -0,0 +1,2 @@
+for i, e in enumerate([]):
+    <caret>

python/testSrc/com/jetbrains/python/postfix/PyEnumEnumeratePostfixTemplateTest.java

@@ -0,0 +1,22 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.postfix;
+
+public class PyEnumEnumeratePostfixTemplateTest extends PyPostfixTemplateTestCase{
+
+  public void testTopLevel() {
+    doTest();
+  }
+
+  public void testFunction() {
+    doTest();
+  }
+
+  public void testComplexExpression() {
+    doTest();
+  }
+
+  @Override
+  protected String getTestDataDir() {
+    return "enum/";
+  }
+}

python/testSrc/com/jetbrains/python/postfix/PyForEnumeratePostfixTemplateTest.java

@@ -0,0 +1,22 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.postfix;
+
+public class PyForEnumeratePostfixTemplateTest extends PyPostfixTemplateTestCase{
+
+  public void testTopLevel() {
+    doTest();
+  }
+
+  public void testFunction() {
+    doTest();
+  }
+
+  public void testComplexExpression() {
+    doTest();
+  }
+
+  @Override
+  protected String getTestDataDir() {
+    return "fore/";
+  }
+}

python/testSrc/com/jetbrains/python/postfix/PyIterEnumeratePostfixTemplateTest.java

@@ -0,0 +1,22 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.postfix;
+
+public class PyIterEnumeratePostfixTemplateTest extends PyPostfixTemplateTestCase{
+
+  public void testTopLevel() {
+    doTest();
+  }
+
+  public void testFunction() {
+    doTest();
+  }
+
+  public void testComplexExpression() {
+    doTest();
+  }
+
+  @Override
+  protected String getTestDataDir() {
+    return "itere/";
+  }
+}