From 807f2a7e032c1d7a8e0aaaa9a1122115286083b5 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Thu, 22 Jun 2017 17:22:33 +0300
Subject: [PATCH] PY-18816 Don't rely on presence of stubs if unstubbing is not
 allowed

Namely, if the stub tree was already replaced by AST, use it instead
but in a restricted manner as if these type aliases were taken from
stubs to prevent stochastic errors (like we do already in other places
in code insight).
---
 .../typing/PyTypingTypeProvider.java          | 22 ++++++++++---------
 .../psi/impl/stubs/PyTypingAliasStubType.java | 12 +++++++---
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
index fd33c8c81f5c..7f41bfc813ba 100644
--- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
@@ -39,6 +39,7 @@ import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyExpressionCodeFragmentImpl;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
+import com.jetbrains.python.psi.impl.stubs.PyTypingAliasStubType;
 import com.jetbrains.python.psi.resolve.*;
 import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
 import com.jetbrains.python.psi.stubs.PyTypingAliasStub;
@@ -786,25 +787,26 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
         // Presumably, a TypeVar definition or a type alias
         if (element instanceof PyTargetExpression) {
           final PyTargetExpression targetExpr = (PyTargetExpression)element;
+          PyExpression assignedValue = null;
           if (context.maySwitchToAST(expression)) {
-            final PyExpression assignedValue = targetExpr.findAssignedValue();
-            if (assignedValue != null) {
-              elements.add(assignedValue);
-              continue;
-            }
+            assignedValue = targetExpr.findAssignedValue();
           }
           else {
             final PyTargetExpressionStub stub = targetExpr.getStub();
             if (stub != null) {
               final PyTypingAliasStub aliasStub = stub.getCustomStub(PyTypingAliasStub.class);
               if (aliasStub != null) {
-                final PyExpression assignedValue = createExpressionFromFragment(aliasStub.getText(), expression);
-                if (assignedValue != null) {
-                  elements.add(assignedValue);
-                  continue;
-                }
+                assignedValue = createExpressionFromFragment(aliasStub.getText(), expression);
               }
             }
+            else {
+              // Use PSI to get the assigned value but only if the same expression would be saved in stubs
+              assignedValue = PyTypingAliasStubType.getAssignedValueIfTypeAliasLike(targetExpr);
+            }
+          }
+          if (assignedValue != null) {
+            elements.add(assignedValue);
+            continue;
           }
         }
         if (isBuiltinPathLike(element)) {
diff --git a/python/src/com/jetbrains/python/psi/impl/stubs/PyTypingAliasStubType.java b/python/src/com/jetbrains/python/psi/impl/stubs/PyTypingAliasStubType.java
index 48fbc75954cd..c6b6d83368c8 100644
--- a/python/src/com/jetbrains/python/psi/impl/stubs/PyTypingAliasStubType.java
+++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyTypingAliasStubType.java
@@ -52,14 +52,20 @@ public class PyTypingAliasStubType extends CustomTargetExpressionStubType<PyTypi
   @Nullable
   @Override
   public PyTypingAliasStub createStub(PyTargetExpression psi) {
-    if (!PyUtil.isTopLevel(psi) || !looksLikeTypeAliasTarget(psi)) {
+    final PyExpression value = getAssignedValueIfTypeAliasLike(psi);
+    return value != null ? new PyTypingTypeAliasStubImpl(value.getText()) : null;
+  }
+
+  @Nullable
+  public static PyExpression getAssignedValueIfTypeAliasLike(@NotNull PyTargetExpression target) {
+    if (!PyUtil.isTopLevel(target) || !looksLikeTypeAliasTarget(target)) {
       return null;
     }
-    final PyExpression value = psi.findAssignedValue();
+    final PyExpression value = target.findAssignedValue();
     if (value == null || !looksLikeTypeHint(value)) {
       return null;
     }
-    return new PyTypingTypeAliasStubImpl(value.getText());
+    return value;
   }
 
   private static boolean looksLikeTypeAliasTarget(@NotNull PyTargetExpression target) {