fix parsing of generator in argument list to actually match Python grammaer (PY-6926)

python/src/com/jetbrains/python/parsing/ExpressionParsing.java

@@ -131,20 +131,26 @@ public class ExpressionParsing extends Parsing {
   }
 
   private void parseComprehension(PsiBuilder.Marker expr,
-                                  final IElementType endToken,
+                                  @Nullable final IElementType endToken,
                                   final IElementType exprType,
                                   final boolean leaveEndTokenOutside) {
     assertCurrentToken(PyTokenTypes.FOR_KEYWORD);
     while (true) {
       myBuilder.advanceLexer();
       parseExpression(true, true);
-      parseComprehensionRange();
+      parseComprehensionRange(exprType == PyElementTypes.GENERATOR_EXPRESSION);
       while (myBuilder.getTokenType() == PyTokenTypes.IF_KEYWORD) {
         myBuilder.advanceLexer();
         if (!parseOldExpression()) {
           myBuilder.error(message("PARSE.expected.expression"));
         }
       }
+      if (atToken(PyTokenTypes.FOR_KEYWORD)) {
+        continue;
+      }
+      if (endToken == null) {
+        break;
+      }
       if (atToken(endToken)) {
         if (leaveEndTokenOutside) {
           expr.done(exprType);
@@ -154,18 +160,22 @@ public class ExpressionParsing extends Parsing {
         nextToken();
         break;
       }
-      if (atToken(PyTokenTypes.FOR_KEYWORD)) {
-        continue;
-      }
       myBuilder.error(message("PARSE.expected.for.or.bracket"));
       break;
     }
     expr.done(exprType);
   }
 
-  protected void parseComprehensionRange() {
+  protected void parseComprehensionRange(boolean generatorExpression) {
     checkMatches(PyTokenTypes.IN_KEYWORD, "'in' expected");
-    if (!parseTupleExpression(false, false, true)) {
+    boolean result;
+    if (generatorExpression) {
+      result = parseORTestExpression(false, false);
+    }
+    else {
+      result = parseTupleExpression(false, false, true);
+    }
+    if (!result) {
       myBuilder.error("expression expected");
     }
   }
@@ -449,20 +459,18 @@ public class ExpressionParsing extends Parsing {
     LOG.assertTrue(myBuilder.getTokenType() == PyTokenTypes.LPAR);
     final PsiBuilder.Marker arglist = myBuilder.mark();
     myBuilder.advanceLexer();
-    final PsiBuilder.Marker genexpr = myBuilder.mark();
+    PsiBuilder.Marker genexpr = myBuilder.mark();
     int argNumber = 0;
-    boolean needBracket = true;
     while (myBuilder.getTokenType() != PyTokenTypes.RPAR) {
       argNumber++;
       if (argNumber > 1) {
-        if (argNumber == 2 && myBuilder.getTokenType() == PyTokenTypes.FOR_KEYWORD && genexpr != null) {
-          parseComprehension(genexpr, PyTokenTypes.RPAR, PyElementTypes.GENERATOR_EXPRESSION, true);
-          needBracket = false;
-          break;
+        if (argNumber == 2 && atToken(PyTokenTypes.FOR_KEYWORD) && genexpr != null) {
+          parseComprehension(genexpr, null, PyElementTypes.GENERATOR_EXPRESSION, true);
+          genexpr = null;
+          continue;
         }
-        else if (myBuilder.getTokenType() == PyTokenTypes.COMMA) {
-          myBuilder.advanceLexer();
-          if (myBuilder.getTokenType() == PyTokenTypes.RPAR) {
+        else if (matchToken(PyTokenTypes.COMMA)) {
+          if (atToken(PyTokenTypes.RPAR)) {
             break;
           }
         }
@@ -500,12 +508,10 @@ public class ExpressionParsing extends Parsing {
       }
     }
 
-    if (needBracket) {
-      if (genexpr != null) {
-        genexpr.drop();
-      }
-      checkMatches(PyTokenTypes.RPAR, message("PARSE.expected.rpar"));
+    if (genexpr != null) {
+      genexpr.drop();
     }
+    checkMatches(PyTokenTypes.RPAR, message("PARSE.expected.rpar"));
     arglist.done(PyElementTypes.ARGUMENT_LIST);
   }
 
@@ -642,7 +648,7 @@ public class ExpressionParsing extends Parsing {
     return true;
   }
 
-  private boolean parseORTestExpression(boolean stopOnIn, boolean isTargetExpression) {
+  protected boolean parseORTestExpression(boolean stopOnIn, boolean isTargetExpression) {
     PsiBuilder.Marker expr = myBuilder.mark();
     if (!parseANDTestExpression(stopOnIn, isTargetExpression)) {
       expr.drop();

python/testData/psi/GeneratorList.py

@@ -0,0 +1 @@
+list(int(i) for i in '1', '2')

python/testData/psi/GeneratorList.txt

@@ -0,0 +1,31 @@
+PyFile:GeneratorList.py
+  PyExpressionStatement
+    PyCallExpression: list
+      PyReferenceExpression: list
+        PsiElement(Py:IDENTIFIER)('list')
+      PyArgumentList
+        PsiElement(Py:LPAR)('(')
+        PyGeneratorExpression
+          PyCallExpression: int
+            PyReferenceExpression: int
+              PsiElement(Py:IDENTIFIER)('int')
+            PyArgumentList
+              PsiElement(Py:LPAR)('(')
+              PyReferenceExpression: i
+                PsiElement(Py:IDENTIFIER)('i')
+              PsiElement(Py:RPAR)(')')
+          PsiWhiteSpace(' ')
+          PsiElement(Py:FOR_KEYWORD)('for')
+          PsiWhiteSpace(' ')
+          PyTargetExpression: i
+            PsiElement(Py:IDENTIFIER)('i')
+          PsiWhiteSpace(' ')
+          PsiElement(Py:IN_KEYWORD)('in')
+          PsiWhiteSpace(' ')
+          PyStringLiteralExpression: 1
+            PsiElement(Py:SINGLE_QUOTED_STRING)(''1'')
+        PsiElement(Py:COMMA)(',')
+        PsiWhiteSpace(' ')
+        PyStringLiteralExpression: 2
+          PsiElement(Py:SINGLE_QUOTED_STRING)(''2'')
+        PsiElement(Py:RPAR)(')')

python/testSrc/com/jetbrains/python/PythonParsingTest.java

@@ -338,6 +338,11 @@ public class PythonParsingTest extends ParsingTestCase {
     doTest();
   }
 
+  // PY-6926
+  public void testGeneratorList() {
+    doTest();
+  }
+
   public void doTest() {
     doTest(LanguageLevel.PYTHON25);
   }