From 83771208a317567780f88a0a1aa1db65307464d3 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Thu, 15 Sep 2016 13:36:24 +0300
Subject: [PATCH] PY-20744 Consider variable annotations only in current file
 to prevent unstabbing everywhere

It will be fixed eventually when we start storing annotations
for parameter, functions and target expressions in stub files and
process them using PyTypeParser only.
---
 .../python/codeInsight/PyTypingTypeProvider.java | 16 +++++++++-------
 .../__init__.py                                  |  0
 .../variableAnnotationsInExternalFiles/lib.py    |  1 +
 .../variableAnnotationsInExternalFiles/main.py   |  3 +++
 .../com/jetbrains/python/PyStubsTest.java        | 13 +++++++++++++
 5 files changed, 26 insertions(+), 7 deletions(-)
 create mode 100644 python/testData/stubs/variableAnnotationsInExternalFiles/__init__.py
 create mode 100644 python/testData/stubs/variableAnnotationsInExternalFiles/lib.py
 create mode 100644 python/testData/stubs/variableAnnotationsInExternalFiles/main.py

diff --git a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
index 3d9a6c001553..cdd550e018fc 100644
--- a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
@@ -176,14 +176,16 @@ 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));
+      if (context.maySwitchToAST(target)) {
+        // 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;
         }
-        return null;
       }
       final String comment = target.getTypeCommentAnnotation();
       if (comment != null) {
diff --git a/python/testData/stubs/variableAnnotationsInExternalFiles/__init__.py b/python/testData/stubs/variableAnnotationsInExternalFiles/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/python/testData/stubs/variableAnnotationsInExternalFiles/lib.py b/python/testData/stubs/variableAnnotationsInExternalFiles/lib.py
new file mode 100644
index 000000000000..80b60df6b81c
--- /dev/null
+++ b/python/testData/stubs/variableAnnotationsInExternalFiles/lib.py
@@ -0,0 +1 @@
+attr: int = ...
\ No newline at end of file
diff --git a/python/testData/stubs/variableAnnotationsInExternalFiles/main.py b/python/testData/stubs/variableAnnotationsInExternalFiles/main.py
new file mode 100644
index 000000000000..662608fcd7b6
--- /dev/null
+++ b/python/testData/stubs/variableAnnotationsInExternalFiles/main.py
@@ -0,0 +1,3 @@
+from .lib import attr
+
+x = attr
diff --git a/python/testSrc/com/jetbrains/python/PyStubsTest.java b/python/testSrc/com/jetbrains/python/PyStubsTest.java
index c4bddf13d69a..a03029b3ee47 100644
--- a/python/testSrc/com/jetbrains/python/PyStubsTest.java
+++ b/python/testSrc/com/jetbrains/python/PyStubsTest.java
@@ -596,4 +596,17 @@ public class PyStubsTest extends PyTestCase {
     assertNotParsed(file1);
     assertNotParsed(file2);
   }
+
+  public void testVariableAnnotationsInExternalFiles() {
+    runWithLanguageLevel(LanguageLevel.PYTHON36, () -> {
+      final PyFile current = getTestFile(getTestName(true) + "/main.py");
+      final PyFile external = getTestFile(getTestName(true) + "/lib.py");
+      final PyTargetExpression attr = current.findTopLevelAttribute("x");
+      assertNotNull(attr);
+      final TypeEvalContext context = TypeEvalContext.codeAnalysis(myFixture.getProject(), current);
+      // Will turn into concrete type when we start saving annotations in stubs 
+      assertNull(context.getType(attr));
+      assertNotParsed(external);
+    });
+  }
 }