diff --git a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java index 38b794fd2d25..f03befac6c14 100644 --- a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java +++ b/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; } diff --git a/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java b/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java index 838a4dbcc75a..f207169df77f 100644 --- a/python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java +++ b/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; diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java index 47d848c73dd7..9513795cb5b7 100644 --- a/python/testSrc/com/jetbrains/python/PyTypingTest.java +++ b/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);