Unify stubs for all Python language levels (PY-26392)

The cause of the problem was the difference in parsing for different
language levels.

This commit improves error recovery and adds support for stubs for
syntax elements unavailable in older Python versions.
Cases that are covered are necessary for stubs generation
for Python 3.6 standard library. They are discovered by automatic
validation executed by stubs generator in StubsGenerator.kt.

Those are the cases that were fixed:
* A stub for single star argument in Py2 was absent
* Type annotations in Py2 were absent
* Print with end argument in Py2 broke parser
* Exec function as argument in Py2 broke parser
* Async keyword after decorator broke parser

All the cases are listed as test cases in PyUnifiedStubsTest

P.S. It is possible though that more cases will be revealed in future.

python/psi-api/src/com/jetbrains/python/psi/LanguageLevel.java

@@ -91,6 +91,11 @@ public enum LanguageLevel {
     return mySupportsSetLiterals;
   }
 
+  public boolean isOutdatedPython2() {
+    return !myIsPy3K;
+  }
+
+  @Deprecated
   public boolean isPy3K() {
     return myIsPy3K;
   }
@@ -152,7 +157,7 @@ public enum LanguageLevel {
   public static LanguageLevel forElement(@NotNull PsiElement element) {
     final PsiFile containingFile = element.getContainingFile();
     if (containingFile instanceof PyFile) {
-      return ((PyFile) containingFile).getLanguageLevel();
+      return ((PyFile)containingFile).getLanguageLevel();
     }
     return getDefault();
   }

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

@@ -37,7 +37,7 @@ public class ExpressionParsing extends Parsing {
 
   public boolean parsePrimaryExpression(boolean isTargetExpression) {
     final IElementType firstToken = myBuilder.getTokenType();
-    if (firstToken == PyTokenTypes.IDENTIFIER) {
+    if (isIdentifier(myBuilder)) {
       if (isTargetExpression) {
         buildTokenElement(PyElementTypes.TARGET_EXPRESSION, myBuilder);
       }
@@ -303,6 +303,20 @@ public class ExpressionParsing extends Parsing {
         parseComprehension(expr, PyTokenTypes.RPAR, PyElementTypes.GENERATOR_EXPRESSION);
       }
       else {
+        final PsiBuilder.Marker err = myBuilder.mark();
+        boolean empty = true;
+        while (myBuilder.getTokenType() != PyTokenTypes.RPAR &&
+               myBuilder.getTokenType() != PyTokenTypes.LINE_BREAK &&
+               myBuilder.getTokenType() != PyTokenTypes.STATEMENT_BREAK) {
+          myBuilder.advanceLexer();
+          empty = false;
+        }
+        if (!empty) {
+          err.error("Unexpected expression syntax");
+        }
+        else {
+          err.drop();
+        }
         checkMatches(PyTokenTypes.RPAR, message("PARSE.expected.rpar"));
         expr.done(PyElementTypes.PARENTHESIZED_EXPRESSION);
       }
@@ -323,7 +337,7 @@ public class ExpressionParsing extends Parsing {
     boolean recastFirstIdentifier = false;
     boolean recastQualifier = false;
     do {
-      boolean firstIdentifierIsTarget = isTargetExpression && ! recastFirstIdentifier;
+      boolean firstIdentifierIsTarget = isTargetExpression && !recastFirstIdentifier;
       PsiBuilder.Marker expr = myBuilder.mark();
       if (!parsePrimaryExpression(firstIdentifierIsTarget)) {
         expr.drop();
@@ -340,7 +354,7 @@ public class ExpressionParsing extends Parsing {
           }
           myBuilder.advanceLexer();
           checkMatches(PyTokenTypes.IDENTIFIER, message("PARSE.expected.name"));
-          if (isTargetExpression && ! recastQualifier && !atAnyOfTokens(PyTokenTypes.DOT, PyTokenTypes.LPAR, PyTokenTypes.LBRACKET)) {
+          if (isTargetExpression && !recastQualifier && !atAnyOfTokens(PyTokenTypes.DOT, PyTokenTypes.LPAR, PyTokenTypes.LBRACKET)) {
             expr.done(PyElementTypes.TARGET_EXPRESSION);
           }
           else {
@@ -390,7 +404,7 @@ public class ExpressionParsing extends Parsing {
               expr.done(PyElementTypes.SUBSCRIPTION_EXPRESSION);
             }
           }
-          if (isTargetExpression && ! recastQualifier) {
+          if (isTargetExpression && !recastQualifier) {
             recastFirstIdentifier = true; // subscription is always a reference
             recastQualifier = true; // recast non-first qualifiers too
             expr.rollbackTo();
@@ -524,9 +538,9 @@ public class ExpressionParsing extends Parsing {
         starArgMarker.done(PyElementTypes.STAR_ARGUMENT_EXPRESSION);
       }
       else {
-        if (myBuilder.getTokenType() == PyTokenTypes.IDENTIFIER) {
+        if (isIdentifier(myBuilder)) {
           final PsiBuilder.Marker keywordArgMarker = myBuilder.mark();
-          myBuilder.advanceLexer();
+          advanceIdentifierLike(myBuilder);
           if (myBuilder.getTokenType() == PyTokenTypes.EQ) {
             myBuilder.advanceLexer();
             if (!parseSingleExpression(false)) {
@@ -637,7 +651,7 @@ public class ExpressionParsing extends Parsing {
       return parseLambdaExpression(false);
     }
     PsiBuilder.Marker condExpr = myBuilder.mark();
-    if (!parseORTestExpression( stopOnIn, isTargetExpression)) {
+    if (!parseORTestExpression(stopOnIn, isTargetExpression)) {
       condExpr.drop();
       return false;
     }
@@ -982,7 +996,8 @@ public class ExpressionParsing extends Parsing {
       else {
         if (isTargetExpression) {
           expr.error("can't assign to await expression");
-        } else {
+        }
+        else {
           expr.done(PyElementTypes.PREFIX_EXPRESSION);
         }
       }
@@ -997,10 +1012,12 @@ public class ExpressionParsing extends Parsing {
   private boolean atForOrAsyncFor() {
     if (atToken(PyTokenTypes.FOR_KEYWORD)) {
       return true;
-    } else if (matchToken(PyTokenTypes.ASYNC_KEYWORD)) {
+    }
+    else if (matchToken(PyTokenTypes.ASYNC_KEYWORD)) {
       if (atToken(PyTokenTypes.FOR_KEYWORD)) {
         return true;
-      } else {
+      }
+      else {
         myBuilder.error("'for' expected");
         return false;
       }

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

@@ -24,6 +24,7 @@ import com.jetbrains.python.psi.LanguageLevel;
 import org.jetbrains.annotations.NotNull;
 
 import static com.jetbrains.python.PyBundle.message;
+import static com.jetbrains.python.parsing.StatementParsing.TOK_ASYNC;
 
 /**
  * @author yole
@@ -57,7 +58,10 @@ public class FunctionParsing extends Parsing {
   }
 
   public void parseReturnTypeAnnotation() {
-    if (myContext.getLanguageLevel().isPy3K() && myBuilder.getTokenType() == PyTokenTypes.RARROW) {
+    if (myBuilder.getTokenType() == PyTokenTypes.RARROW) {
+      if (myContext.getLanguageLevel().isOutdatedPython2()) {
+        myBuilder.error("Return type annotations are unsupported in Python 2");
+      }
       PsiBuilder.Marker maybeReturnAnnotation = myBuilder.mark();
       nextToken();
       if (!myContext.getExpressionParser().parseSingleExpression(false)) {
@@ -104,6 +108,10 @@ public class FunctionParsing extends Parsing {
       myBuilder.advanceLexer();
       parseDeclarationAfterDecorator(endMarker, true);
     }
+    else if (atToken(PyTokenTypes.IDENTIFIER, TOK_ASYNC)) {
+      advanceAsync(true);
+      parseDeclarationAfterDecorator(endMarker, true);
+    }
     else {
       parseDeclarationAfterDecorator(endMarker, false);
     }
@@ -186,15 +194,20 @@ public class FunctionParsing extends Parsing {
   }
 
   protected boolean parseParameter(IElementType endToken, boolean isLambda) {
-    final PsiBuilder.Marker parameter = myBuilder.mark();
+    PsiBuilder.Marker parameter = myBuilder.mark();
     boolean isStarParameter = false;
     if (myBuilder.getTokenType() == PyTokenTypes.MULT) {
       myBuilder.advanceLexer();
-      if (myContext.getLanguageLevel().isPy3K() &&
-          (myBuilder.getTokenType() == PyTokenTypes.COMMA) || myBuilder.getTokenType() == endToken) {
+      if ((myBuilder.getTokenType() == PyTokenTypes.COMMA) || myBuilder.getTokenType() == endToken) {
+        if (myContext.getLanguageLevel().isOutdatedPython2()) {
+          parameter.rollbackTo();
+          parameter = myBuilder.mark();
+          advanceError(myBuilder, "Single star parameter is not supported in Python 2");
+        }
         parameter.done(PyElementTypes.SINGLE_STAR_PARAMETER);
         return true;
       }
+
       isStarParameter = true;
     }
     else if (myBuilder.getTokenType() == PyTokenTypes.EXP) {
@@ -208,7 +221,7 @@ public class FunctionParsing extends Parsing {
       if (!isStarParameter && matchToken(PyTokenTypes.EQ)) {
         if (!getExpressionParser().parseSingleExpression(false)) {
           PsiBuilder.Marker invalidElements = myBuilder.mark();
-          while(!atAnyOfTokens(endToken, PyTokenTypes.LINE_BREAK, PyTokenTypes.COMMA, null)) {
+          while (!atAnyOfTokens(endToken, PyTokenTypes.LINE_BREAK, PyTokenTypes.COMMA, null)) {
             nextToken();
           }
           invalidElements.error(message("PARSE.expected.expression"));
@@ -232,7 +245,10 @@ public class FunctionParsing extends Parsing {
   }
 
   public void parseParameterAnnotation() {
-    if (myContext.getLanguageLevel().isPy3K() && atToken(PyTokenTypes.COLON)) {
+    if (atToken(PyTokenTypes.COLON)) {
+      if (myContext.getLanguageLevel().isOutdatedPython2()) {
+        myBuilder.error("Type annotations are unsupported in Python 2");
+      }
       PsiBuilder.Marker annotationMarker = myBuilder.mark();
       nextToken();
       if (!getExpressionParser().parseSingleExpression(false)) {

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

@@ -110,9 +110,43 @@ public class Parsing {
     myBuilder.advanceLexer();
   }
 
+  protected void advanceAsync(boolean falseAsync) {
+    if (falseAsync) {
+      advanceError(myBuilder, "'async' keyword is not expected here");
+    }
+    else {
+      myBuilder.advanceLexer();
+    }
+  }
+
+  protected static void advanceIdentifierLike(@NotNull PsiBuilder builder) {
+    if (isFalseIdentifier(builder)) {
+      String tokenText = builder.getTokenText();
+      advanceError(builder, "'" + tokenText + "' keyword can't be used as identifier in Python 2");
+    }
+    else {
+      builder.advanceLexer();
+    }
+  }
+
+  protected static void advanceError(@NotNull PsiBuilder builder, String message) {
+    final PsiBuilder.Marker err = builder.mark();
+    builder.advanceLexer();
+    err.error(message);
+  }
+
+  protected static boolean isIdentifier(@NotNull PsiBuilder builder) {
+    return builder.getTokenType() == PyTokenTypes.IDENTIFIER || isFalseIdentifier(builder);
+  }
+
+  private static boolean isFalseIdentifier(@NotNull PsiBuilder builder) {
+    return builder.getTokenType() == PyTokenTypes.EXEC_KEYWORD ||
+           builder.getTokenType() == PyTokenTypes.PRINT_KEYWORD;
+  }
+
   protected static void buildTokenElement(IElementType type, PsiBuilder builder) {
     final PsiBuilder.Marker marker = builder.mark();
-    builder.advanceLexer();
+    advanceIdentifierLike(builder);
     marker.done(type);
   }
 

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

@@ -50,7 +50,7 @@ public class StatementParsing extends Parsing implements ITokenTypeRemapper {
   @NonNls protected static final String TOK_FALSE = "False";
   @NonNls protected static final String TOK_NONLOCAL = "nonlocal";
   @NonNls protected static final String TOK_EXEC = "exec";
-  @NonNls protected static final String TOK_ASYNC = "async";
+  @NonNls public static final String TOK_ASYNC = "async";
   @NonNls protected static final String TOK_AWAIT = "await";
 
   private static final String EXPRESSION_EXPECTED = "Expression expected";
@@ -133,10 +133,16 @@ public class StatementParsing extends Parsing implements ITokenTypeRemapper {
       return;
     }
     if (firstToken == PyTokenTypes.ASYNC_KEYWORD) {
-      parseAsyncStatement();
+      parseAsyncStatement(false);
       return;
     }
 
+    if (atToken(PyTokenTypes.IDENTIFIER, TOK_ASYNC)) {
+      if (parseAsyncStatement(true)) {
+        return;
+      }
+    }
+
     parseSimpleStatement();
   }
 
@@ -655,7 +661,7 @@ public class StatementParsing extends Parsing implements ITokenTypeRemapper {
     while (!atAnyOfTokens(null, PyTokenTypes.DEDENT, PyTokenTypes.STATEMENT_BREAK, PyTokenTypes.COLON)) {
       myBuilder.advanceLexer();
     }
-    boolean result =  matchToken(PyTokenTypes.COLON);
+    boolean result = matchToken(PyTokenTypes.COLON);
     if (!result && atToken(PyTokenTypes.STATEMENT_BREAK)) {
       myBuilder.advanceLexer();
     }
@@ -821,23 +827,36 @@ public class StatementParsing extends Parsing implements ITokenTypeRemapper {
     classMarker.done(PyElementTypes.CLASS_DECLARATION);
   }
 
-  private void parseAsyncStatement() {
-    assertCurrentToken(PyTokenTypes.ASYNC_KEYWORD);
+  private boolean parseAsyncStatement(boolean falseAsync) {
+    if (!falseAsync) {
+      assertCurrentToken(PyTokenTypes.ASYNC_KEYWORD);
+    }
     final PsiBuilder.Marker marker = myBuilder.mark();
-    myBuilder.advanceLexer();
+
+    advanceAsync(falseAsync);
+
     final IElementType token = myBuilder.getTokenType();
     if (token == PyTokenTypes.DEF_KEYWORD) {
       getFunctionParser().parseFunctionDeclaration(marker, true);
+      return true;
     }
     else if (token == PyTokenTypes.WITH_KEYWORD) {
       parseWithStatement(marker);
+      return true;
     }
     else if (token == PyTokenTypes.FOR_KEYWORD) {
       parseForStatement(marker);
+      return true;
     }
     else {
-      marker.drop();
-      myBuilder.error("'def' or 'with' or 'for' expected");
+      if (falseAsync) {
+        marker.rollbackTo();
+      }
+      else {
+        marker.drop();
+        myBuilder.error("'def' or 'with' or 'for' expected");
+      }
+      return false;
     }
   }
 

python/testData/highlighting/async.py

@@ -22,7 +22,7 @@ async = 1
         <info descr="null">async</info> for x in xs:
             pass
 
-    async<error descr="End of statement expected"> </error>with <info descr="PY.PARAMETER">xs</info>:
+    <error descr="'async' keyword is not expected here">async</error> with <info descr="PY.PARAMETER">xs</info>:
         pass
 
     return async

python/testData/psi/AsyncDefWithDecorator.py

@@ -0,0 +1,4 @@
+class AsyncIterator(AsyncIterable):
+    @abstractmethod
+    async def __anext__(self):
+        raise StopAsyncIteration

python/testData/psi/AsyncFor.txt

@@ -30,13 +30,10 @@ PyFile:AsyncFor.py
             PyPassStatement
               PsiElement(Py:PASS_KEYWORD)('pass')
   PsiWhiteSpace('\n\n\n')
-  PyExpressionStatement
-    PyReferenceExpression: async
-      PsiElement(Py:IDENTIFIER)('async')
-    PsiErrorElement:End of statement expected
-      <empty list>
-  PsiWhiteSpace(' ')
   PyForStatement
+    PsiErrorElement:'async' keyword is not expected here
+      PsiElement(Py:IDENTIFIER)('async')
+    PsiWhiteSpace(' ')
     PyForPart
       PsiElement(Py:FOR_KEYWORD)('for')
       PsiWhiteSpace(' ')

python/testData/psi/AsyncWith.txt

@@ -25,13 +25,10 @@ PyFile:AsyncWith.py
           PyPassStatement
             PsiElement(Py:PASS_KEYWORD)('pass')
   PsiWhiteSpace('\n\n')
-  PyExpressionStatement
-    PyReferenceExpression: async
-      PsiElement(Py:IDENTIFIER)('async')
-    PsiErrorElement:End of statement expected
-      <empty list>
-  PsiWhiteSpace(' ')
   PyWithStatement
+    PsiErrorElement:'async' keyword is not expected here
+      PsiElement(Py:IDENTIFIER)('async')
+    PsiWhiteSpace(' ')
     PsiElement(Py:WITH_KEYWORD)('with')
     PsiWhiteSpace(' ')
     PyWithItem

python/testData/psi/ExecAsArgument.py

@@ -0,0 +1,6 @@
+def foo():
+    baz(exec, 1)
+
+
+def bar():
+    pass

python/testData/psi/PrintWithEndArgument.py

@@ -0,0 +1,6 @@
+def bar():
+    print("header:", end=" ")
+
+
+def foo():
+    pass

python/testData/psi/StarParameter.py

@@ -0,0 +1,6 @@
+def foo(arg1, *, kwarg1):
+    pass
+
+
+def bar():
+    pass

python/testSrc/com/jetbrains/python/PyUnifiedStubsTest.kt

@@ -15,30 +15,95 @@
  */
 package com.jetbrains.python
 
-import com.intellij.psi.PsiManager
+import com.intellij.psi.util.PsiTreeUtil
 import com.jetbrains.python.fixtures.PyTestCase
 import com.jetbrains.python.psi.LanguageLevel
 import com.jetbrains.python.psi.PyFile
+import com.jetbrains.python.psi.PySingleStarParameter
 import junit.framework.TestCase
 
 /**
  * @author traff
  */
 
+/**
+ * Stubs should be language level independent.
+ *
+ * This test checks the cases when stubs can differ for different language levels.
+ *
+ */
 class PyUnifiedStubsTest : PyTestCase() {
 
   override fun getTestDataPath(): String {
     return PythonTestUtil.getTestDataPath()
   }
 
-  fun testExecPy3() {
+  private fun doTest(function: (languageLevel: LanguageLevel, file: PyFile) -> Boolean) {
     for (level in LanguageLevel.SUPPORTED_LEVELS) {
       runWithLanguageLevel(level) {
-        val vf = myFixture.copyFileToProject("psi/${getTestName(false)}.py")
-        val file = PsiManager.getInstance(myFixture.project).findFile(vf) as PyFile
+        val file = getFile()
 
-        TestCase.assertEquals("Python $level", "exec", file.topLevelFunctions[0].name)
+        val shouldContainErrors = function(level, file)
+
+        if (shouldContainErrors) {
+          TestCase.assertTrue(PsiTreeUtil.hasErrorElements(file))
+        }
       }
+      tearDown() // re-init fixture to avoid using cached psi
+      setUp()
     }
   }
+
+  private fun getFile(): PyFile {
+    val vf = myFixture.copyFileToProject("psi/${getTestName(false)}.py")
+    val file = myFixture.psiManager.findFile(vf) as PyFile
+    return file
+  }
+
+  fun testExecPy3() {
+    doTest { level, file ->
+      TestCase.assertEquals("Python $level", "exec", file.topLevelFunctions[0].name)
+      level.isOutdatedPython2
+    }
+  }
+
+  fun testStarParameter() {
+    doTest { level, file ->
+      val functionStub = file.topLevelFunctions[0].stub
+      TestCase.assertNotNull("Function should contain star argument stub",
+                             functionStub!!.childrenStubs[0].findChildStubByType<PySingleStarParameter>(
+                               PyElementTypes.SINGLE_STAR_PARAMETER))
+      level.isOutdatedPython2
+    }
+  }
+
+  fun testAnnotations() {
+    doTest { level, file ->
+      val functionStub = file.topLevelFunctions[0].stub
+      TestCase.assertNotNull("Function should contain annotation stub", functionStub!!.findChildStubByType(PyElementTypes.ANNOTATION))
+      level.isOutdatedPython2
+    }
+  }
+
+  fun testPrintWithEndArgument() {
+    doTest { level, file ->
+      TestCase.assertEquals("'end' argument shouldn't break tuple parsing for Python 2", 2, file.topLevelFunctions.size)
+      level.isOutdatedPython2
+    }
+  }
+
+  fun testExecAsArgument() {
+    doTest { level, file ->
+      TestCase.assertEquals("'exec' used as an argument shouldn't break parser for Python 2", 2, file.topLevelFunctions.size)
+      level.isOutdatedPython2
+    }
+  }
+
+  fun testAsyncDefWithDecorator() {
+    doTest { level, file ->
+      TestCase.assertEquals("Async keyword shouldn't break function parsing for Python <3.5", 1, file.topLevelClasses[0].methods.size)
+      level.isOlderThan(LanguageLevel.PYTHON35)
+    }
+  }
+
 }

python/tools/src/com/jetbrains/python/tools/PyPrebuiltIndicesGenerator.kt

@@ -15,6 +15,7 @@ import com.intellij.util.ui.UIUtil
 import com.jetbrains.python.psi.impl.stubs.PyPrebuiltStubsProvider
 import org.jetbrains.index.id.IdIndexGenerator
 import java.io.File
+import java.util.*
 
 fun main(args: Array<String>) {
   buildIndices(args[0], "${args[1]}/${PyPrebuiltStubsProvider.NAME}")

tools/index-tools/src/org/jetbrains/index/stubs/StubsGenerator.kt

@@ -36,7 +36,6 @@ import com.intellij.util.indexing.FileContentImpl
 import com.intellij.util.io.PersistentHashMap
 import junit.framework.TestCase
 import org.jetbrains.index.IndexGenerator
-import java.io.ByteArrayInputStream
 import java.io.File
 import java.util.*
 
@@ -226,21 +225,24 @@ abstract class LanguageLevelAwareStubsGenerator<T>(stubsVersion: String, stubsSt
 
       applyLanguageLevel(languageLevel)
 
+      // create new FileContentImpl, because it caches stub in user data
+      val content = FileContentImpl(fileContent.file, fileContent.content)
 
-      val stub = super.buildStubForFile(fileContent, serializationManager)
+      val stub = super.buildStubForFile(content, serializationManager)
 
       val bytes = BufferExposingByteArrayOutputStream()
       serializationManager.serialize(stub, bytes)
 
-      if (prevLanguageLevelBytes != null) {
-        if (!Arrays.equals(bytes.toByteArray(), prevLanguageLevelBytes)) {
-          val stub2 = serializationManager.deserialize(ByteArrayInputStream(prevLanguageLevelBytes))
-          val msg = "Stubs are different for ${fileContent.file.path} between Python versions $prevLanguageLevel and $languageLevel.\n"
-          TestCase.assertEquals(msg, DebugUtil.stubTreeToString(stub), DebugUtil.stubTreeToString(stub2))
+      if (prevStub != null) {
+        try {
+          check(stub, prevStub)
+        } catch (e: AssertionError) {
+          val msg = "Stubs are different for ${content.file.path} between Python versions $prevLanguageLevel and $languageLevel.\n"
+          TestCase.assertEquals(msg, DebugUtil.stubTreeToString(stub), DebugUtil.stubTreeToString(prevStub))
           TestCase.fail(msg + "But DebugUtil.stubTreeToString values of stubs are unfortunately equal.")
         }
       }
-      prevLanguageLevelBytes = bytes.toByteArray()
+
       prevLanguageLevel = languageLevel
       prevStub = stub
     }
@@ -251,3 +253,16 @@ abstract class LanguageLevelAwareStubsGenerator<T>(stubsVersion: String, stubsSt
   }
 }
 
+private fun check(stub: Stub, stub2: Stub) {
+  TestCase.assertEquals(stub.stubType, stub2.stubType)
+  val stubs = stub.childrenStubs
+  val stubs2 = stub2.childrenStubs
+  TestCase.assertEquals(stubs.size, stubs2.size)
+  var i = 0
+  val len = stubs.size
+  while (i < len) {
+    check(stubs[i], stubs2[i])
+    ++i
+  }
+}
+