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.

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);

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;
+  }
 }

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')

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);