From 6d86b858ca3e31d85c7a5951eb4470e170009c96 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Wed, 14 Sep 2016 15:34:56 +0300
Subject: [PATCH] =?UTF-8?q?PY-20744=20Initial=20support=20of=20variable=20?=
 =?UTF-8?q?annotations=20in=20type=20=D1=81hecking?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../python/psi/PyTargetExpression.java          |  2 +-
 .../codeInsight/PyTypingTypeProvider.java       |  9 +++++++++
 .../python/psi/impl/PyTargetExpressionImpl.java | 10 ++++++++++
 .../com/jetbrains/python/PyTypingTest.java      | 17 ++++++++++++++++-
 4 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyTargetExpression.java b/python/psi-api/src/com/jetbrains/python/psi/PyTargetExpression.java
index 9012137bb53f..a248dee2d144 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/PyTargetExpression.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyTargetExpression.java
@@ -27,7 +27,7 @@ import org.jetbrains.annotations.Nullable;
  */
 public interface PyTargetExpression extends PyQualifiedExpression, PsiNamedElement, PsiNameIdentifierOwner, PyDocStringOwner,
                                             PyQualifiedNameOwner, PyReferenceOwner, StubBasedPsiElement<PyTargetExpressionStub>,
-                                            PyPossibleClassMember, PyTypeCommentOwner {
+                                            PyPossibleClassMember, PyTypeCommentOwner, PyAnnotationOwner {
   PyTargetExpression[] EMPTY_ARRAY = new PyTargetExpression[0];
 
   /**
diff --git a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
index bf017fc484c0..3d9a6c001553 100644
--- a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
@@ -176,6 +176,15 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
   public PyType getReferenceType(@NotNull PsiElement referenceTarget, TypeEvalContext context, @Nullable PsiElement anchor) {
     if (referenceTarget instanceof PyTargetExpression) {
       final PyTargetExpression target = (PyTargetExpression)referenceTarget;
+      // XXX: Requires switching from stub to AST
+      final PyAnnotation annotation = target.getAnnotation();
+      if (annotation != null) {
+        final PyExpression value = annotation.getValue();
+        if (value != null) {
+          return getType(value, new Context(context));
+        }
+        return null;
+      }
       final String comment = target.getTypeCommentAnnotation();
       if (comment != null) {
         final PyType type = getStringBasedType(comment, referenceTarget, new Context(context));
diff --git a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
index 868370081cc7..90e30082cf6e 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
@@ -214,6 +214,16 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
     }
   }
 
+  @Nullable
+  @Override
+  public PyAnnotation getAnnotation() {
+    final PsiElement parent = getParent();
+    if (parent instanceof PyAssignmentStatement || parent instanceof PyTypeDeclarationStatement) {
+      return ((PyAnnotationOwner)parent).getAnnotation();
+    }
+    return null;
+  }
+
   @Nullable
   private static PyType getWithItemVariableType(TypeEvalContext context, PyWithItem item) {
     final PyExpression expression = item.getExpression();
diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java
index 9c60998664e6..3ef21a127898 100644
--- a/python/testSrc/com/jetbrains/python/PyTypingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java
@@ -711,7 +711,22 @@ public class PyTypingTest extends PyTestCase {
            "    sublist = x[0]\n" +
            "    expr = sublist[0]\n");
   }
-  
+
+  public void testLocalVariableAnnotation() {
+    doTest("int",
+           "def f():\n" +
+           "    x: int = undefined()\n" +
+           "    expr = x");
+  }
+
+  public void testInstanceAttributeAnnotation() {
+    doTest("int",
+           "class C:\n" +
+           "    attr: int\n" +
+           "    \n" +
+           "expr = C().attr");
+  }
+
   private void doTestNoInjectedText(@NotNull String text) {
     myFixture.configureByText(PythonFileType.INSTANCE, text);
     final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myFixture.getProject());