PY-20744 Parse PEP-526 variable annotations

Annotation is preserved at the level of assignment nodes similar to
where CPython keeps them in its AST (in special "augassign" nodes).
For type annotations in form "x: int" without variable initialization
special statement PyTypeDefinitionStatement was introduced.

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyAnnotationOwner.java b/python/psi-api/src/com/jetbrains/python/psi/PyAnnotationOwner.java
new file mode 100644 (file)
index 0000000..952e83e
--- /dev/null
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyAnnotationOwner.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2000-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.psi;
+
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * @author Mikhail Golubev
+ */
+public interface PyAnnotationOwner {
+  @Nullable
+  PyAnnotation getAnnotation();
+}

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyAssignmentStatement.java b/python/psi-api/src/com/jetbrains/python/psi/PyAssignmentStatement.java
index b4b2a0470060c807a57f1a8c511e719132afbe02..38c68735e7da6378f2401cb3e4833b2acabe49b9 100644 (file)
--- a/python/psi-api/src/com/jetbrains/python/psi/PyAssignmentStatement.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyAssignmentStatement.java
@@ -24,7 +24,7 @@ import java.util.List;
 /**
  * Describes an assignment statement.
  */
-public interface PyAssignmentStatement extends PyStatement, PyNamedElementContainer {
+public interface PyAssignmentStatement extends PyStatement, PyNamedElementContainer, PyAnnotationOwner {
 
   /**
    * @return the left-hand side of the statement; each item may consist of many elements.

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyElementVisitor.java b/python/psi-api/src/com/jetbrains/python/psi/PyElementVisitor.java
index c6c900ab0a5bfcd228b1bbc73a9af0ebbb856fc9..076dd026db10ac67ee450834a8a734582d21f687 100644 (file)
--- a/python/psi-api/src/com/jetbrains/python/psi/PyElementVisitor.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyElementVisitor.java
@@ -280,4 +280,8 @@ public class PyElementVisitor extends PsiElementVisitor {
   public void visitPyWithItem(PyWithItem node) {
     visitPyElement(node);
   }
+
+  public void visitPyTypeDeclarationStatement(PyTypeDeclarationStatement node) {
+    visitPyStatement(node);
+  }
 }

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java b/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
index cad8ad22612a4a7f5a90aeefd412f7b59f02dacd..25d6a20bb691354a191e785fb15c8a080b12a332 100644 (file)
--- a/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
@@ -36,7 +36,7 @@ import java.util.List;
  */
 public interface PyFunction extends PsiNamedElement, StubBasedPsiElement<PyFunctionStub>, PsiNameIdentifierOwner, PyStatement, PyCallable,
                                     PyDocStringOwner, ScopeOwner, PyDecoratable, PyTypedElement, PyStatementListContainer,
-                                    PyPossibleClassMember, PyTypeCommentOwner {
+                                    PyPossibleClassMember, PyTypeCommentOwner, PyAnnotationOwner {
 
   PyFunction[] EMPTY_ARRAY = new PyFunction[0];
   ArrayFactory<PyFunction> ARRAY_FACTORY = count -> new PyFunction[count];

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyNamedParameter.java b/python/psi-api/src/com/jetbrains/python/psi/PyNamedParameter.java
index dcb268af254d3237ac1e05910ef568468999439f..48b5e6ec473ba648b8f50f78f9eb0c639d625a49 100644 (file)
--- a/python/psi-api/src/com/jetbrains/python/psi/PyNamedParameter.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyNamedParameter.java
@@ -20,12 +20,12 @@ import com.intellij.psi.PsiNamedElement;
 import com.intellij.psi.StubBasedPsiElement;
 import com.jetbrains.python.psi.stubs.PyNamedParameterStub;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 
 /**
  * Represents a named parameter, as opposed to a tuple parameter.
  */
-public interface PyNamedParameter extends PyParameter, PsiNamedElement, PsiNameIdentifierOwner, PyExpression, PyTypeCommentOwner, StubBasedPsiElement<PyNamedParameterStub> {
+public interface PyNamedParameter extends PyParameter, PsiNamedElement, PsiNameIdentifierOwner, PyExpression, PyTypeCommentOwner,
+                                          PyAnnotationOwner, StubBasedPsiElement<PyNamedParameterStub> {
   boolean isPositionalContainer();
 
   boolean isKeywordContainer();
@@ -44,8 +44,5 @@ public interface PyNamedParameter extends PyParameter, PsiNamedElement, PsiNameI
    */
   @NotNull
   String getRepr(boolean includeDefaultValue);
-
-  @Nullable
-  PyAnnotation getAnnotation();
 }
 

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyTypeDeclarationStatement.java b/python/psi-api/src/com/jetbrains/python/psi/PyTypeDeclarationStatement.java
new file mode 100644 (file)
index 0000000..299cc0d
--- /dev/null
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyTypeDeclarationStatement.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2000-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.psi;
+
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * @author Mikhail Golubev
+ */
+public interface PyTypeDeclarationStatement extends PyStatement, PyAnnotationOwner {
+  @NotNull
+  PyExpression getTarget();
+}

diff --git a/python/src/com/jetbrains/python/PyElementTypes.java b/python/src/com/jetbrains/python/PyElementTypes.java
index 30c4ef84269434245227e91c89ff60ea88eac9a4..09f7bd64178e903abf5cadcd7d64d718420a0f5b 100644 (file)
--- a/python/src/com/jetbrains/python/PyElementTypes.java
+++ b/python/src/com/jetbrains/python/PyElementTypes.java
@@ -62,6 +62,7 @@ public interface PyElementTypes {
   PyElementType DEL_STATEMENT = new PyElementType("DEL_STATEMENT", PyDelStatementImpl.class);
   PyElementType EXEC_STATEMENT = new PyElementType("EXEC_STATEMENT", PyExecStatementImpl.class);
   PyElementType FOR_STATEMENT = new PyElementType("FOR_STATEMENT", PyForStatementImpl.class);
+  PyElementType TYPE_DECLARATION_STATEMENT = new PyElementType("TYPE_DECLARATION_STATEMENT", PyTypeDeclarationStatementImpl.class); 
 
   PyStubElementType<PyFromImportStatementStub, PyFromImportStatement> FROM_IMPORT_STATEMENT = new PyFromImportStatementElementType();
   PyStubElementType<PyImportStatementStub, PyImportStatement> IMPORT_STATEMENT = new PyImportStatementElementType();

diff --git a/python/src/com/jetbrains/python/PyTypeDeclarationStatementImpl.java b/python/src/com/jetbrains/python/PyTypeDeclarationStatementImpl.java
new file mode 100644 (file)
index 0000000..85e7938
--- /dev/null
+++ b/python/src/com/jetbrains/python/PyTypeDeclarationStatementImpl.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright 2000-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python;
+
+import com.intellij.lang.ASTNode;
+import com.jetbrains.python.psi.PyAnnotation;
+import com.jetbrains.python.psi.PyElementVisitor;
+import com.jetbrains.python.psi.PyExpression;
+import com.jetbrains.python.psi.PyTypeDeclarationStatement;
+import com.jetbrains.python.psi.impl.PyElementImpl;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * @author Mikhail Golubev
+ */
+public class PyTypeDeclarationStatementImpl extends PyElementImpl implements PyTypeDeclarationStatement {
+  public PyTypeDeclarationStatementImpl(ASTNode astNode) {
+    super(astNode);
+  }
+
+  @NotNull
+  @Override
+  public PyExpression getTarget() {
+    return findNotNullChildByClass(PyExpression.class);
+  }
+
+  @Nullable
+  @Override
+  public PyAnnotation getAnnotation() {
+    return findChildByClass(PyAnnotation.class);
+  }
+
+  @Override
+  protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
+    pyVisitor.visitPyTypeDeclarationStatement(this);
+  }
+}

diff --git a/python/src/com/jetbrains/python/PythonTokenSetContributor.java b/python/src/com/jetbrains/python/PythonTokenSetContributor.java
index 5e78c9ac6743532b2a4ecb35a3a6438a3bdb0ba4..63d432a694eea4d55575ac0b0c71632ddcd8d23c 100644 (file)
--- a/python/src/com/jetbrains/python/PythonTokenSetContributor.java
+++ b/python/src/com/jetbrains/python/PythonTokenSetContributor.java
@@ -28,8 +28,8 @@ public class PythonTokenSetContributor extends PythonDialectsTokenSetContributor
   @NotNull
   @Override
   public TokenSet getStatementTokens() {
-    return TokenSet.create(EXPRESSION_STATEMENT, ASSIGNMENT_STATEMENT, AUG_ASSIGNMENT_STATEMENT, ASSERT_STATEMENT,
-                           BREAK_STATEMENT, CONTINUE_STATEMENT, DEL_STATEMENT, EXEC_STATEMENT, FOR_STATEMENT,
+    return TokenSet.create(EXPRESSION_STATEMENT, ASSIGNMENT_STATEMENT, TYPE_DECLARATION_STATEMENT, AUG_ASSIGNMENT_STATEMENT, 
+                           ASSERT_STATEMENT, BREAK_STATEMENT, CONTINUE_STATEMENT, DEL_STATEMENT, EXEC_STATEMENT, FOR_STATEMENT,
                            FROM_IMPORT_STATEMENT, GLOBAL_STATEMENT, IMPORT_STATEMENT, IF_STATEMENT, PASS_STATEMENT,
                            PRINT_STATEMENT, RAISE_STATEMENT, RETURN_STATEMENT, TRY_EXCEPT_STATEMENT, WITH_STATEMENT,
                            WHILE_STATEMENT, NONLOCAL_STATEMENT, CLASS_DECLARATION, FUNCTION_DECLARATION);

diff --git a/python/src/com/jetbrains/python/parsing/FunctionParsing.java b/python/src/com/jetbrains/python/parsing/FunctionParsing.java
index 05d359ebc0d7ba6d84842aad0fcac6c8d6762073..519242a45c05be14b71c1a461b581cce796c26c4 100644 (file)
--- a/python/src/com/jetbrains/python/parsing/FunctionParsing.java
+++ b/python/src/com/jetbrains/python/parsing/FunctionParsing.java
@@ -230,7 +230,7 @@ public class FunctionParsing extends Parsing {
     return true;
   }
 
-  protected void parseParameterAnnotation() {
+  public void parseParameterAnnotation() {
     if (myContext.getLanguageLevel().isPy3K() && atToken(PyTokenTypes.COLON)) {
       PsiBuilder.Marker annotationMarker = myBuilder.mark();
       nextToken();

diff --git a/python/src/com/jetbrains/python/parsing/StatementParsing.java b/python/src/com/jetbrains/python/parsing/StatementParsing.java
index 70c229262811f1901ab3e0e7d49ad4263ba5d783..ec54d4755b5e9fd31d8b50100b8b60dfa7a7787d 100644 (file)
--- a/python/src/com/jetbrains/python/parsing/StatementParsing.java
+++ b/python/src/com/jetbrains/python/parsing/StatementParsing.java
@@ -214,37 +214,45 @@ public class StatementParsing extends Parsing implements ITokenTypeRemapper {
           builder.error(EXPRESSION_EXPECTED);
         }
       }
-      else if (builder.getTokenType() == PyTokenTypes.EQ) {
-        statementType = PyElementTypes.ASSIGNMENT_STATEMENT;
+      else if (atToken(PyTokenTypes.EQ) || (atToken(PyTokenTypes.COLON) && myContext.getLanguageLevel().isPy3K())) {
         exprStatement.rollbackTo();
         exprStatement = builder.mark();
         getExpressionParser().parseExpression(false, true);
-        LOG.assertTrue(builder.getTokenType() == PyTokenTypes.EQ, builder.getTokenType());
-        builder.advanceLexer();
+        LOG.assertTrue(builder.getTokenType() == PyTokenTypes.EQ || builder.getTokenType() == PyTokenTypes.COLON, builder.getTokenType());
 
-        while (true) {
-          PsiBuilder.Marker maybeExprMarker = builder.mark();
-          final boolean isYieldExpr = builder.getTokenType() == PyTokenTypes.YIELD_KEYWORD;
-          if (!getExpressionParser().parseYieldOrTupleExpression(false)) {
-            maybeExprMarker.drop();
-            builder.error(EXPRESSION_EXPECTED);
-            break;
-          }
-          if (builder.getTokenType() == PyTokenTypes.EQ) {
-            if (isYieldExpr) {
+        if (builder.getTokenType() == PyTokenTypes.COLON) {
+          statementType = PyElementTypes.TYPE_DECLARATION_STATEMENT;
+          getFunctionParser().parseParameterAnnotation();
+        }
+
+        if (builder.getTokenType() == PyTokenTypes.EQ) {
+          statementType = PyElementTypes.ASSIGNMENT_STATEMENT;
+          builder.advanceLexer();
+
+          while (true) {
+            PsiBuilder.Marker maybeExprMarker = builder.mark();
+            final boolean isYieldExpr = builder.getTokenType() == PyTokenTypes.YIELD_KEYWORD;
+            if (!getExpressionParser().parseYieldOrTupleExpression(false)) {
               maybeExprMarker.drop();
-              builder.error("Cannot assign to 'yield' expression");
+              builder.error(EXPRESSION_EXPECTED);
+              break;
+            }
+            if (builder.getTokenType() == PyTokenTypes.EQ) {
+              if (isYieldExpr) {
+                maybeExprMarker.drop();
+                builder.error("Cannot assign to 'yield' expression");
+              }
+              else {
+                maybeExprMarker.rollbackTo();
+                getExpressionParser().parseExpression(false, true);
+                LOG.assertTrue(builder.getTokenType() == PyTokenTypes.EQ, builder.getTokenType());
+              }
+              builder.advanceLexer();
             }
             else {
-              maybeExprMarker.rollbackTo();
-              getExpressionParser().parseExpression(false, true);
-              LOG.assertTrue(builder.getTokenType() == PyTokenTypes.EQ, builder.getTokenType());
+              maybeExprMarker.drop();
+              break;
             }
-            builder.advanceLexer();
-          }
-          else {
-            maybeExprMarker.drop();
-            break;
           }
         }
       }

diff --git a/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java
index c05082490f679b47e326ad0c54bfa4356e1015f3..08281fc811a33b90eb0c6a6cef2f4bb36078df61 100644 (file)
--- a/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java
@@ -96,6 +96,12 @@ public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssign
     return targets.toArray(new PyExpression[targets.size()]);
   }
 
+  @Nullable
+  @Override
+  public PyAnnotation getAnnotation() {
+    return findChildByClass(PyAnnotation.class);
+  }
+
   private static void addCandidate(List<PyExpression> candidates, PyExpression psi) {
     if (psi instanceof PyParenthesizedExpression) {
       addCandidate(candidates, ((PyParenthesizedExpression)psi).getContainedExpression());

diff --git a/python/testData/psi/VariableAnnotations.py b/python/testData/psi/VariableAnnotations.py
new file mode 100644 (file)
index 0000000..0b6800b
--- /dev/null
+++ b/python/testData/psi/VariableAnnotations.py
@@ -0,0 +1,8 @@
+class C:
+    x: int
+    y: None = 42
+
+    def m(self, d):
+        x: List[bool]
+        d['foo']: str
+        (d['bar']): float

diff --git a/python/testData/psi/VariableAnnotations.txt b/python/testData/psi/VariableAnnotations.txt
new file mode 100644 (file)
index 0000000..3b847b9
--- /dev/null
+++ b/python/testData/psi/VariableAnnotations.txt
@@ -0,0 +1,93 @@
+PyFile:VariableAnnotations.py
+  PyClass: C
+    PsiElement(Py:CLASS_KEYWORD)('class')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:IDENTIFIER)('C')
+    PyArgumentList
+      <empty list>
+    PsiElement(Py:COLON)(':')
+    PsiWhiteSpace('\n    ')
+    PyStatementList
+      PyTypeDeclarationStatement
+        PyTargetExpression: x
+          PsiElement(Py:IDENTIFIER)('x')
+        PyAnnotation
+          PsiElement(Py:COLON)(':')
+          PsiWhiteSpace(' ')
+          PyReferenceExpression: int
+            PsiElement(Py:IDENTIFIER)('int')
+      PsiWhiteSpace('\n    ')
+      PyAssignmentStatement
+        PyTargetExpression: y
+          PsiElement(Py:IDENTIFIER)('y')
+        PyAnnotation
+          PsiElement(Py:COLON)(':')
+          PsiWhiteSpace(' ')
+          PyNoneLiteralExpression
+            PsiElement(Py:NONE_KEYWORD)('None')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:EQ)('=')
+        PsiWhiteSpace(' ')
+        PyNumericLiteralExpression
+          PsiElement(Py:INTEGER_LITERAL)('42')
+      PsiWhiteSpace('\n\n    ')
+      PyFunction('m')
+        PsiElement(Py:DEF_KEYWORD)('def')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:IDENTIFIER)('m')
+        PyParameterList
+          PsiElement(Py:LPAR)('(')
+          PyNamedParameter('self')
+            PsiElement(Py:IDENTIFIER)('self')
+          PsiElement(Py:COMMA)(',')
+          PsiWhiteSpace(' ')
+          PyNamedParameter('d')
+            PsiElement(Py:IDENTIFIER)('d')
+          PsiElement(Py:RPAR)(')')
+        PsiElement(Py:COLON)(':')
+        PsiWhiteSpace('\n        ')
+        PyStatementList
+          PyTypeDeclarationStatement
+            PyTargetExpression: x
+              PsiElement(Py:IDENTIFIER)('x')
+            PyAnnotation
+              PsiElement(Py:COLON)(':')
+              PsiWhiteSpace(' ')
+              PySubscriptionExpression
+                PyReferenceExpression: List
+                  PsiElement(Py:IDENTIFIER)('List')
+                PsiElement(Py:LBRACKET)('[')
+                PyReferenceExpression: bool
+                  PsiElement(Py:IDENTIFIER)('bool')
+                PsiElement(Py:RBRACKET)(']')
+          PsiWhiteSpace('\n        ')
+          PyTypeDeclarationStatement
+            PySubscriptionExpression
+              PyReferenceExpression: d
+                PsiElement(Py:IDENTIFIER)('d')
+              PsiElement(Py:LBRACKET)('[')
+              PyStringLiteralExpression: foo
+                PsiElement(Py:SINGLE_QUOTED_STRING)(''foo'')
+              PsiElement(Py:RBRACKET)(']')
+            PyAnnotation
+              PsiElement(Py:COLON)(':')
+              PsiWhiteSpace(' ')
+              PyReferenceExpression: str
+                PsiElement(Py:IDENTIFIER)('str')
+          PsiWhiteSpace('\n        ')
+          PyTypeDeclarationStatement
+            PyParenthesizedExpression
+              PsiElement(Py:LPAR)('(')
+              PySubscriptionExpression
+                PyReferenceExpression: d
+                  PsiElement(Py:IDENTIFIER)('d')
+                PsiElement(Py:LBRACKET)('[')
+                PyStringLiteralExpression: bar
+                  PsiElement(Py:SINGLE_QUOTED_STRING)(''bar'')
+                PsiElement(Py:RBRACKET)(']')
+              PsiElement(Py:RPAR)(')')
+            PyAnnotation
+              PsiElement(Py:COLON)(':')
+              PsiWhiteSpace(' ')
+              PyReferenceExpression: float
+                PsiElement(Py:IDENTIFIER)('float')
\ No newline at end of file

diff --git a/python/testSrc/com/jetbrains/python/PythonParsingTest.java b/python/testSrc/com/jetbrains/python/PythonParsingTest.java
index 169c0141848b6eba4590855bf545092219b8a306..e90e63aff9f7fa19812e449b49c8d2307caf4538 100644 (file)
--- a/python/testSrc/com/jetbrains/python/PythonParsingTest.java
+++ b/python/testSrc/com/jetbrains/python/PythonParsingTest.java
@@ -515,6 +515,10 @@ public class PythonParsingTest extends ParsingTestCase {
     doTest(LanguageLevel.PYTHON35);
   }
 
+  public void testVariableAnnotations() {
+    doTest(LanguageLevel.PYTHON36);
+  }
+
   public void doTest(LanguageLevel languageLevel) {
     LanguageLevel prev = myLanguageLevel;
     myLanguageLevel = languageLevel;