Support for string-based 'typing' type annotations

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

@@ -17,11 +17,13 @@ package com.jetbrains.python.codeInsight;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiPolyVariantReference;
 import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyExpressionCodeFragmentImpl;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.types.*;
@@ -158,6 +160,27 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
     if (functionType != null) {
       return functionType;
     }
+    final PyType stringBasedType = getStringBasedType(expression, context);
+    if (stringBasedType != null) {
+      return stringBasedType;
+    }
+    return null;
+  }
+
+  @Nullable
+  private static PyType getStringBasedType(@NotNull PyExpression expression, @NotNull TypeEvalContext context) {
+    if (expression instanceof PyStringLiteralExpression) {
+      // XXX: Requires switching from stub to AST
+      final String contents = ((PyStringLiteralExpression)expression).getStringValue();
+      final Project project = expression.getProject();
+      final PyExpressionCodeFragmentImpl codeFragment = new PyExpressionCodeFragmentImpl(project, "dummy.py", contents, false);
+      codeFragment.setContext(expression.getContainingFile());
+      final PsiElement element = codeFragment.getFirstChild();
+      if (element instanceof PyExpressionStatement) {
+        final PyExpression dummyExpr = ((PyExpressionStatement)element).getExpression();
+        return getType(dummyExpr, context);
+      }
+    }
     return null;
   }
 

python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java

@@ -101,7 +101,13 @@ public class PyResolveUtil {
                                   @Nullable PsiElement roof) {
     // Use real context here to enable correct completion and resolve in case of PyExpressionCodeFragment!!!
     final PsiElement realContext = PyPsiUtils.getRealContext(element);
-    final ScopeOwner originalOwner = ScopeUtil.getScopeOwner(realContext);
+    final ScopeOwner originalOwner;
+    if (realContext != element && realContext instanceof PyFile) {
+      originalOwner = (PyFile)realContext;
+    }
+    else {
+      originalOwner = ScopeUtil.getScopeOwner(realContext);
+    }
     final PsiElement parent = element.getParent();
     final boolean isGlobalOrNonlocal = parent instanceof PyGlobalStatement || parent instanceof PyNonlocalStatement;
     ScopeOwner owner = originalOwner;

python/testSrc/com/jetbrains/python/PyTypingTest.java

@@ -232,6 +232,23 @@ public class PyTypingTest extends PyTestCase {
            "    pass\n");
   }
 
+  public void testTypeInStringLiteral() {
+    doTest("C",
+           "class C:\n" +
+           "    def foo(self, expr: 'C'):\n" +
+           "        pass\n");
+  }
+
+  public void testQualifiedTypeInStringLiteral() {
+    doTest("str",
+           "import typing\n" +
+           "\n" +
+           "def foo(x: 'typing.AnyStr') -> typing.AnyStr:\n" +
+           "    pass\n" +
+           "\n" +
+           "expr = foo('bar')\n");
+  }
+
   private void doTest(@NotNull String expectedType, @NotNull String text) {
     myFixture.copyDirectoryToProject("typing", "");
     myFixture.configureByText(PythonFileType.INSTANCE, text);