From f20ab38e84953d027abb8c4b2b3f9776d6ad72da Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Fri, 23 Jun 2017 16:03:55 +0300
Subject: [PATCH] PY-18816 Exclude qualified references from type aliases stubs

Since, otherwise, these custom stubs conflict with the standard
"initializer" field of PyTargetExpressionStub and, thus, break
resolve in stubbed files, e.g. when exported symbols are aliased in
"__init__.py" of a package, etc. The current workaround is not to
keep RHS text of such assignments in the custom stubs, relying
on existing functionality of PyTypingAliasStubType instead, but inspect
both when extracting type aliases from the stub tree.
---
 .../typing/PyTypingTypeProvider.java          | 18 ++-------
 .../psi/impl/stubs/PyTypingAliasStubType.java | 40 ++++++++++++++++---
 python/testData/stubs/TypeAliasStubs.py       |  5 ++-
 .../com/jetbrains/python/PyStubsTest.java     |  5 +++
 4 files changed, 47 insertions(+), 21 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
index 7f41bfc813ba..f9b895986b6c 100644
--- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
@@ -41,8 +41,6 @@ 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;
 import com.jetbrains.python.psi.types.*;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
@@ -615,7 +613,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
   }
 
   @Nullable
-  private static PyExpression createExpressionFromFragment(@NotNull String contents, @NotNull PsiElement anchor) {
+  public static PyExpression createExpressionFromFragment(@NotNull String contents, @NotNull PsiElement anchor) {
     final PyExpressionCodeFragmentImpl codeFragment = new PyExpressionCodeFragmentImpl(anchor.getProject(), "dummy.py", contents, false);
     codeFragment.setContext(FileContextUtil.getContextFile(anchor));
     final PyExpressionStatement statement = as(codeFragment.getFirstChild(), PyExpressionStatement.class);
@@ -787,22 +785,12 @@ 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;
+          final PyExpression assignedValue;
           if (context.maySwitchToAST(expression)) {
             assignedValue = targetExpr.findAssignedValue();
           }
           else {
-            final PyTargetExpressionStub stub = targetExpr.getStub();
-            if (stub != null) {
-              final PyTypingAliasStub aliasStub = stub.getCustomStub(PyTypingAliasStub.class);
-              if (aliasStub != null) {
-                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);
-            }
+            assignedValue = PyTypingAliasStubType.getAssignedValueStubSafe(targetExpr);
           }
           if (assignedValue != null) {
             elements.add(assignedValue);
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 c6b6d83368c8..ea59fb65b7e7 100644
--- a/python/src/com/jetbrains/python/psi/impl/stubs/PyTypingAliasStubType.java
+++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyTypingAliasStubType.java
@@ -23,11 +23,14 @@ import com.intellij.util.io.StringRef;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
+import com.jetbrains.python.psi.stubs.PyTargetExpressionStub.InitializerType;
 import com.jetbrains.python.psi.stubs.PyTypingAliasStub;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.io.IOException;
+import java.util.Objects;
 import java.util.regex.Pattern;
 
 import static com.jetbrains.python.psi.PyUtil.as;
@@ -52,17 +55,17 @@ public class PyTypingAliasStubType extends CustomTargetExpressionStubType<PyTypi
   @Nullable
   @Override
   public PyTypingAliasStub createStub(PyTargetExpression psi) {
-    final PyExpression value = getAssignedValueIfTypeAliasLike(psi);
+    final PyExpression value = getAssignedValueIfTypeAliasLike(psi, true);
     return value != null ? new PyTypingTypeAliasStubImpl(value.getText()) : null;
   }
 
   @Nullable
-  public static PyExpression getAssignedValueIfTypeAliasLike(@NotNull PyTargetExpression target) {
+  private static PyExpression getAssignedValueIfTypeAliasLike(@NotNull PyTargetExpression target, boolean forStubCreation) {
     if (!PyUtil.isTopLevel(target) || !looksLikeTypeAliasTarget(target)) {
       return null;
     }
     final PyExpression value = target.findAssignedValue();
-    if (value == null || !looksLikeTypeHint(value)) {
+    if (value == null || !looksLikeTypeHint(value, forStubCreation)) {
       return null;
     }
     return value;
@@ -84,7 +87,7 @@ public class PyTypingAliasStubType extends CustomTargetExpressionStubType<PyTypi
     return targets.length == 1 && targets[0] == target;
   }
 
-  private static boolean looksLikeTypeHint(@NotNull PyExpression expression) {
+  private static boolean looksLikeTypeHint(@NotNull PyExpression expression, boolean forStubCreation) {
     final PyCallExpression call = as(expression, PyCallExpression.class);
     if (call != null) {
       final PyReferenceExpression callee = as(call.getCallee(), PyReferenceExpression.class);
@@ -100,7 +103,10 @@ public class PyTypingAliasStubType extends CustomTargetExpressionStubType<PyTypi
       return TYPE_ANNOTATION_LIKE.matcher(content).matches();
     }
 
-    if (expression instanceof PyReferenceExpression || expression instanceof PySubscriptionExpression) {
+    // Plain reference expressions are handled by PyTargetExpressionStub.getInitializer() 
+    // when initializer type is ReferenceExpression. We don't want to override these stubs,
+    // so as not to break existing resolve functionality (see PyTargetExpression.getAssignedQName()).
+    if ((!forStubCreation && expression instanceof PyReferenceExpression) || expression instanceof PySubscriptionExpression) {
       return isSyntacticallyValidAnnotation(expression);
     }
     
@@ -129,4 +135,28 @@ public class PyTypingAliasStubType extends CustomTargetExpressionStubType<PyTypi
     final StringRef ref = stream.readName();
     return ref != null ? new PyTypingTypeAliasStubImpl(ref.getString()) : null;
   }
+
+  @Nullable
+  public static PyExpression getAssignedValueStubSafe(@NotNull PyTargetExpression target) {
+    final PyTargetExpressionStub stub = target.getStub();
+    PyExpression result = null;
+    if (stub != null) {
+      final PyTypingAliasStub aliasStub = stub.getCustomStub(PyTypingAliasStub.class);
+      String aliasText = null;
+      if (aliasStub != null) {
+        aliasText = aliasStub.getText();
+      }
+      else if (stub.getInitializerType() == InitializerType.ReferenceExpression) {
+        aliasText = Objects.toString(stub.getInitializer(), null);
+      }
+      if (aliasText != null) {
+        result = PyTypingTypeProvider.createExpressionFromFragment(aliasText, target);
+      }
+    }
+    else {
+      // Use PSI to get the assigned value but only if the same expression would be saved in stubs
+      result = getAssignedValueIfTypeAliasLike(target, false);
+    }
+    return result;
+  }
 }
diff --git a/python/testData/stubs/TypeAliasStubs.py b/python/testData/stubs/TypeAliasStubs.py
index acc30b9fdae3..f0a4c31e9a7d 100644
--- a/python/testData/stubs/TypeAliasStubs.py
+++ b/python/testData/stubs/TypeAliasStubs.py
@@ -6,7 +6,10 @@ S1_ok = "foo"
 S2_ok = "foo.bar"
 S3_ok = "foo.bar[baz]"
 
-plain_ref_ok = foo.bar.baz
+# Such expressions are kept as qualified expressions in PyTargetExpressionStub 
+# with initializer type of ReferenceExpression instead of custom stubs for
+# typing aliases
+plain_ref = foo.bar.baz
 illegal_ref = foo[42].bar.baz
 
 T1_ok = TypeVar('T1')
diff --git a/python/testSrc/com/jetbrains/python/PyStubsTest.java b/python/testSrc/com/jetbrains/python/PyStubsTest.java
index 964826f2ee46..2d7416e4921b 100644
--- a/python/testSrc/com/jetbrains/python/PyStubsTest.java
+++ b/python/testSrc/com/jetbrains/python/PyStubsTest.java
@@ -783,6 +783,11 @@ public class PyStubsTest extends PyTestCase {
       assertHasTypingAliasStub(attr.getName().endsWith("_ok"), attr);
     }
 
+    final PyTargetExpression referenceAlias = file.findTopLevelAttribute("plain_ref");
+    final PyTargetExpressionStub referenceAliasStub = referenceAlias.getStub();
+    assertEquals(PyTargetExpressionStub.InitializerType.ReferenceExpression, referenceAliasStub.getInitializerType());
+    assertEquals(QualifiedName.fromDottedString("foo.bar.baz"), referenceAliasStub.getInitializer());
+
     final PyClass pyClass = file.findTopLevelClass("C");
     final TypeEvalContext context = TypeEvalContext.codeInsightFallback(myFixture.getProject());
     final PyTargetExpression classAttr = pyClass.findClassAttribute("class_attr", false, context);