PY-18816 Store RHS text of assignments looking like type aliases in stubs

python/psi-api/src/com/jetbrains/python/psi/stubs/PyTypingAliasStub.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2000-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.psi.stubs;
+
+import com.jetbrains.python.psi.impl.stubs.CustomTargetExpressionStub;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * @author Mikhail Golubev
+ */
+public interface PyTypingAliasStub extends CustomTargetExpressionStub {
+  @NotNull
+  String getText();
+}

python/src/META-INF/python-core-common.xml

@@ -678,6 +678,7 @@
 
     <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PropertyStubType"/>
     <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PyNamedTupleStubType"/>
+    <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PyTypingAliasStubType"/>
 
     <dialectsTokenSetContributor implementation="com.jetbrains.python.PythonTokenSetContributor"/>
     <pyClassMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibClassMembersProvider"/>

python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java

@@ -40,6 +40,8 @@ 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.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;
@@ -761,8 +763,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
   private static List<PsiElement> tryResolving(@NotNull PyExpression expression, @NotNull TypeEvalContext context) {
     final List<PsiElement> elements = Lists.newArrayList();
     if (expression instanceof PyReferenceExpression) {
-      final PyReferenceExpression referenceExpr = (PyReferenceExpression)expression;
-      final List<PsiElement> results = tryResolvingOnStubs(referenceExpr, context);
+      final List<PsiElement> results = tryResolvingOnStubs((PyReferenceExpression)expression, context);
       for (PsiElement element : results) {
         if (element instanceof PyFunction) {
           final PyFunction function = (PyFunction)element;
@@ -785,11 +786,25 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
         // Presumably, a TypeVar definition or a type alias
         if (element instanceof PyTargetExpression) {
           final PyTargetExpression targetExpr = (PyTargetExpression)element;
-          // XXX: Requires switching from stub to AST
-          final PyExpression assignedValue = targetExpr.findAssignedValue();
-          if (assignedValue != null) {
-            elements.add(assignedValue);
-            continue;
+          if (context.maySwitchToAST(expression)) {
+            final PyExpression assignedValue = targetExpr.findAssignedValue();
+            if (assignedValue != null) {
+              elements.add(assignedValue);
+              continue;
+            }
+          }
+          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;
+                }
+              }
+            }
           }
         }
         if (isBuiltinPathLike(element)) {
@@ -811,23 +826,23 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
   }
 
   @NotNull
-  private static List<PsiElement> tryResolvingOnStubs(@NotNull PyExpression expression,
+  private static List<PsiElement> tryResolvingOnStubs(@NotNull PyReferenceExpression expression,
                                                       @NotNull TypeEvalContext context) {
     
     // PyPsiUtils.asQualifiedName() also takes into account subscription and prefix expressions
-    final QualifiedName qualifiedName = makeQualifiedNameFromReferenceExpression(expression);
+    final QualifiedName qualifiedName = turnPlainReferenceExpressionIntoQualifiedName(expression);
     final PyFile pyFile = as(FileContextUtil.getContextFile(expression), PyFile.class);
 
     if (pyFile != null && qualifiedName != null && qualifiedName.getComponentCount() > 0) {
       List<RatedResolveResult> results = new ArrayList<>();
       //noinspection ConstantConditions
       results.addAll(pyFile.multiResolveName(qualifiedName.getFirstComponent(), false));
-      if (results.isEmpty() && expression instanceof PyQualifiedExpression) {
+      if (results.isEmpty()) {
         for (PyReferenceResolveProvider provider : Extensions.getExtensions(PyReferenceResolveProvider.EP_NAME)) {
           if (provider instanceof PyOverridingReferenceResolveProvider) {
             continue;
           }
-          results.addAll(provider.resolveName((PyQualifiedExpression)expression, context));
+          results.addAll(provider.resolveName(expression, context));
         }
       }
 
@@ -855,8 +870,19 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
     return Collections.singletonList(expression);
   }
 
+  /**
+   * Return the qualified name containing all names in the given (possibly qualified) reference expression.
+   * If any of the qualifiers is not a reference expression, returns null.
+   * <p>
+   * For instance, for the expression "foo.bar.baz" it returns the qualified name "foo.bar.baz",
+   * but for "foo[0].bar.baz" it will return null.
+   * <p>
+   * If you need to take into account such implicit "magical" names, use {@link com.jetbrains.python.psi.impl.PyPsiUtils#asQualifiedName(PyExpression)}
+   * or {@link PyQualifiedExpression#asQualifiedName()}.
+   * @param expression
+   */
   @Nullable
-  private static QualifiedName makeQualifiedNameFromReferenceExpression(@NotNull PyExpression expression) {
+  public static QualifiedName turnPlainReferenceExpressionIntoQualifiedName(@NotNull PyReferenceExpression expression) {
     final List<String> components = new ArrayList<>();
     PyExpression remaining = expression;
     while (remaining != null) {

python/src/com/jetbrains/python/psi/PyFileElementType.java

@@ -62,7 +62,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
   @Override
   public int getStubVersion() {
     // Don't forget to update versions of indexes that use the updated stub-based elements
-    return 61;
+    return 62;
   }
 
   @Nullable

python/src/com/jetbrains/python/psi/impl/stubs/PyTypingAliasStubType.java

@@ -0,0 +1,126 @@
+/*
+ * Copyright 2000-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.psi.impl.stubs;
+
+import com.intellij.extapi.psi.ASTDelegatePsiElement;
+import com.intellij.psi.stubs.StubInputStream;
+import com.intellij.psi.tree.TokenSet;
+import com.intellij.psi.util.PsiTreeUtil;
+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.PyTypingAliasStub;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.io.IOException;
+import java.util.regex.Pattern;
+
+import static com.jetbrains.python.psi.PyUtil.as;
+
+/**
+ * @author Mikhail Golubev
+ */
+public class PyTypingAliasStubType extends CustomTargetExpressionStubType<PyTypingAliasStub> {
+  private static final int STRING_LITERAL_LENGTH_THRESHOLD = 120;
+
+  private static final Pattern TYPE_ANNOTATION_LIKE = Pattern.compile("\\p{javaJavaIdentifierStart}\\p{javaJavaIdentifierPart}*" +
+                                                                      "(\\.\\p{javaJavaIdentifierStart}\\p{javaJavaIdentifierPart}*)*" +
+                                                                      "(\\[.*])?$");
+
+  private static final TokenSet VALID_TYPE_ANNOTATION_ELEMENTS = TokenSet.create(PyElementTypes.REFERENCE_EXPRESSION,
+                                                                                 PyElementTypes.SUBSCRIPTION_EXPRESSION,
+                                                                                 PyElementTypes.TUPLE_EXPRESSION, 
+                                                                                 // List of types is allowed only inside Callable[...]
+                                                                                 PyElementTypes.LIST_LITERAL_EXPRESSION,
+                                                                                 PyElementTypes.STRING_LITERAL_EXPRESSION);
+
+  @Nullable
+  @Override
+  public PyTypingAliasStub createStub(PyTargetExpression psi) {
+    if (!PyUtil.isTopLevel(psi) || !looksLikeTypeAliasTarget(psi)) {
+      return null;
+    }
+    final PyExpression value = psi.findAssignedValue();
+    if (value == null || !looksLikeTypeHint(value)) {
+      return null;
+    }
+    return new PyTypingTypeAliasStubImpl(value.getText());
+  }
+
+  private static boolean looksLikeTypeAliasTarget(@NotNull PyTargetExpression target) {
+    if (target.isQualified()) {
+      return false;
+    }
+    final String name = target.getName();
+    if (name == null || PyUtil.isSpecialName(name)) {
+      return false;
+    }
+    final PyAssignmentStatement assignment = PsiTreeUtil.getParentOfType(target, PyAssignmentStatement.class);
+    if (assignment == null) {
+      return false;
+    }
+    final PyExpression[] targets = assignment.getRawTargets();
+    return targets.length == 1 && targets[0] == target;
+  }
+
+  private static boolean looksLikeTypeHint(@NotNull PyExpression expression) {
+    final PyCallExpression call = as(expression, PyCallExpression.class);
+    if (call != null) {
+      final PyReferenceExpression callee = as(call.getCallee(), PyReferenceExpression.class);
+      return callee != null && "TypeVar".equals(callee.getReferencedName());
+    }
+
+    final PyStringLiteralExpression pyString = as(expression, PyStringLiteralExpression.class);
+    if (pyString != null) {
+      if (pyString.getStringNodes().size() != 1 && pyString.getTextLength() > STRING_LITERAL_LENGTH_THRESHOLD) {
+        return false;
+      }
+      final String content = pyString.getStringValue();
+      return TYPE_ANNOTATION_LIKE.matcher(content).matches();
+    }
+
+    if (expression instanceof PyReferenceExpression || expression instanceof PySubscriptionExpression) {
+      return isSyntacticallyValidAnnotation(expression);
+    }
+    
+    return false;
+  }
+
+  private static boolean isSyntacticallyValidAnnotation(@NotNull PyExpression expression) {
+    return PsiTreeUtil.processElements(expression, element -> {
+      // Check only composite elements
+      if (element instanceof ASTDelegatePsiElement) {
+        if (!VALID_TYPE_ANNOTATION_ELEMENTS.contains(element.getNode().getElementType())) {
+          return false;
+        }
+        if (element instanceof PyReferenceExpression) {
+          // too complex reference expression, e.g. foo[bar].baz
+          return PyTypingTypeProvider.turnPlainReferenceExpressionIntoQualifiedName((PyReferenceExpression)element) != null;
+        }
+      }
+      return true;
+    });
+  }
+
+  @Nullable
+  @Override
+  public PyTypingAliasStub deserializeStub(StubInputStream stream) throws IOException {
+    final StringRef ref = stream.readName();
+    return ref != null ? new PyTypingTypeAliasStubImpl(ref.getString()) : null;
+  }
+}

python/src/com/jetbrains/python/psi/impl/stubs/PyTypingTypeAliasStubImpl.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2000-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.psi.impl.stubs;
+
+import com.intellij.psi.stubs.StubOutputStream;
+import com.intellij.psi.util.QualifiedName;
+import com.jetbrains.python.psi.stubs.PyTypingAliasStub;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.io.IOException;
+
+/**
+ * @author Mikhail Golubev
+ */
+public class PyTypingTypeAliasStubImpl implements PyTypingAliasStub {
+  private String myText;
+
+  public PyTypingTypeAliasStubImpl(@NotNull String text) {
+    myText = text;
+  }
+
+  @NotNull
+  @Override
+  public String getText() {
+    return myText;
+  }
+
+  @NotNull
+  @Override
+  public Class<? extends CustomTargetExpressionStubType> getTypeClass() {
+    return PyTypingAliasStubType.class;
+  }
+
+  @Override
+  public void serialize(StubOutputStream stream) throws IOException {
+    stream.writeName(myText);
+  }
+
+  @Nullable
+  @Override
+  public QualifiedName getCalleeName() {
+    return null;
+  }
+}

python/testData/stubs/TypeAliasInParameterAnnotation.py

@@ -0,0 +1,7 @@
+from typing import Dict, Any
+
+JsonObject = Dict[str, Any]
+
+
+def func(x: JsonObject):
+    pass

python/testData/stubs/TypeAliasStubs.py

@@ -0,0 +1,38 @@
+__author__ = 'Mikhail.Golubev'
+__all__ = ['S1', 'S2']
+__version__ = '0.1'
+
+S1_ok = "foo"
+S2_ok = "foo.bar"
+S3_ok = "foo.bar[baz]"
+
+plain_ref_ok = foo.bar.baz
+illegal_ref = foo[42].bar.baz
+
+T1_ok = TypeVar('T1')
+T2_ok = typing.TypeVar('T2')
+T3 = func()
+
+global_list = [1, 2, 3]
+global_tuple = (1, 2, 3)
+
+for for_counter in range(10):
+    pass
+
+xs_comp = [comp_counter for comp_counter in range(10)]
+
+multi_assign1 = multi_assign2 = Any
+unpack1, unpack2 = Any
+complex.ref = Any
+
+illegal_generic1 = table[table[0].foo]
+illegal_generic2 = Dict[(str, int)]
+illegal_generic3 = Tuple[int, 3]
+illegal_generic4 = Optional[func()]
+
+
+class C:
+    class_attr = Any
+
+    def __init__(self):
+        self.inst_attr = Any

python/testSrc/com/jetbrains/python/PyStubsTest.java

@@ -35,10 +35,7 @@ import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyFileImpl;
 import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher;
-import com.jetbrains.python.psi.stubs.PyClassNameIndex;
-import com.jetbrains.python.psi.stubs.PyNamedTupleStub;
-import com.jetbrains.python.psi.stubs.PySuperClassIndex;
-import com.jetbrains.python.psi.stubs.PyVariableNameIndex;
+import com.jetbrains.python.psi.stubs.*;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import com.jetbrains.python.toolbox.Maybe;
@@ -766,4 +763,50 @@ public class PyStubsTest extends PyTestCase {
       assertNotParsed(file);
     });
   }
+
+  // PY-18116
+  public void testTypeAliasInParameterAnnotation() {
+    runWithLanguageLevel(LanguageLevel.PYTHON30, () -> {
+      final PyFile file = getTestFile();
+      final PyFunction func = file.findTopLevelFunction("func");
+      final PyNamedParameter param = func.getParameterList().findParameterByName("x");
+      assertType("Dict[str, Any]", param, TypeEvalContext.codeInsightFallback(myFixture.getProject()));
+      assertNotParsed(file);
+    });
+  }
+
+  // PY-18116
+  public void testTypeAliasStubs() {
+    final PyFile file = getTestFile();
+    final List<PyTargetExpression> attributes = file.getTopLevelAttributes();
+    for (PyTargetExpression attr : attributes) {
+      assertHasTypingAliasStub(attr.getName().endsWith("_ok"), attr);
+    }
+
+    final PyClass pyClass = file.findTopLevelClass("C");
+    final TypeEvalContext context = TypeEvalContext.codeInsightFallback(myFixture.getProject());
+    final PyTargetExpression classAttr = pyClass.findClassAttribute("class_attr", false, context);
+    assertHasTypingAliasStub(false, classAttr);
+
+    final PyTargetExpression instanceAttr = pyClass.findInstanceAttribute("inst_attr", false);
+    assertHasTypingAliasStub(false, instanceAttr);
+    assertNotParsed(file);
+  }
+
+  @Nullable
+  private static PyTypingAliasStub getAliasStub(@NotNull PyTargetExpression targetExpression) {
+    final PyTargetExpressionStub stub = targetExpression.getStub();
+    return stub != null ? stub.getCustomStub(PyTypingAliasStub.class) : null;
+  }
+
+  private static void assertHasTypingAliasStub(boolean has, @NotNull PyTargetExpression expression) {
+    final String message = "Target '" + expression.getName() + "' should " + (has ? "" : "not ") + "have typing alias stub";
+    final PyTypingAliasStub stub = getAliasStub(expression);
+    if (has) {
+      assertNotNull(message, stub);
+    }
+    else {
+      assertNull(message, stub);
+    }
+  }
 }

python/testSrc/com/jetbrains/python/PythonCompletionTest.java

@@ -22,11 +22,15 @@ import com.intellij.codeInsight.lookup.AutoCompletionPolicy;
 import com.intellij.codeInsight.lookup.Lookup;
 import com.intellij.codeInsight.lookup.LookupElement;
 import com.intellij.codeInsight.lookup.LookupElementBuilder;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.psi.util.QualifiedName;
 import com.intellij.testFramework.PsiTestUtil;
 import com.jetbrains.python.documentation.PyDocumentationSettings;
 import com.jetbrains.python.documentation.docstrings.DocStringFormat;
 import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.LanguageLevel;
+import com.jetbrains.python.psi.PyReferenceExpression;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;