From d2eb0c6fde8c2c789318942a781d738d4349b624 Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Sat, 25 Nov 2017 20:43:05 +0300 Subject: [PATCH] PY-21864 Support of ahead of time variable annotations We search for them only in the same scope where the corresponding target expression (of for loop, with statement or LHS of assignment with tuple unpacking) is defined itself. Another limitation is that this lookup is performed without proper caching of this information on an affected target yet. --- .../typing/PyTypingTypeProvider.java | 74 ++++++++++++------- .../com/jetbrains/python/PyTypingTest.java | 34 +++++++++ 2 files changed, 81 insertions(+), 27 deletions(-) diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java index a94e4c6710f1..cdd1132c14ce 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java @@ -19,6 +19,9 @@ import com.intellij.util.containers.HashMap; import com.intellij.util.containers.HashSet; import com.jetbrains.python.PyCustomType; import com.jetbrains.python.PyNames; +import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache; +import com.jetbrains.python.codeInsight.controlflow.ScopeOwner; +import com.jetbrains.python.codeInsight.dataflow.scope.Scope; import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil; import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyFunctionTypeAnnotation; import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyFunctionTypeAnnotationFile; @@ -332,39 +335,56 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { return annotatedType.get(); } - final PyClass pyClass = target.getContainingClass(); - final PyFunction method = as(ScopeUtil.getScopeOwner(target), PyFunction.class); - if (pyClass != null && method != null && target.isQualified()) { - final String name = target.getReferencedName(); - final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context); + final String name = target.getReferencedName(); + final ScopeOwner scopeOwner = ScopeUtil.getScopeOwner(target); + if (name == null || scopeOwner == null) { + return null; + } - boolean isInstanceAttribute = false; - if (context.maySwitchToAST(target)) { - isInstanceAttribute = StreamEx.of(PyUtil.multiResolveTopPriority(target.getQualifier(), resolveContext)) - .select(PyParameter.class) - .filter(PyParameter::isSelf) - .anyMatch(p -> PsiTreeUtil.getParentOfType(p, PyFunction.class) == method); + final PyClass pyClass = target.getContainingClass(); + + if (target.isQualified()) { + if (pyClass != null && scopeOwner instanceof PyFunction) { + final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context); + + boolean isInstanceAttribute = false; + if (context.maySwitchToAST(target)) { + isInstanceAttribute = StreamEx.of(PyUtil.multiResolveTopPriority(target.getQualifier(), resolveContext)) + .select(PyParameter.class) + .filter(PyParameter::isSelf) + .anyMatch(p -> PsiTreeUtil.getParentOfType(p, PyFunction.class) == scopeOwner); + } + else { + isInstanceAttribute = PyUtil.isInstanceAttribute(target); + } + if (!isInstanceAttribute) { + return null; + } + // Set isDefinition=true to start searching right from the class level. + final PyClassTypeImpl classType = new PyClassTypeImpl(pyClass, true); + final List classAttrs = classType.resolveMember(name, target, AccessDirection.READ, resolveContext, true); + if (classAttrs == null) { + return null; + } + return StreamEx.of(classAttrs) + .map(RatedResolveResult::getElement) + .select(PyTargetExpression.class) + .filter(x -> ScopeUtil.getScopeOwner(x) instanceof PyClass) + .map(x -> getTypeFromTargetExpressionAnnotation(x, context)) + .nonNull() + .map(Ref::get) + .foldLeft(PyUnionType::union) + .orElse(null); } - else { - isInstanceAttribute = PyUtil.isInstanceAttribute(target); - } - if (!isInstanceAttribute) { - return null; - } - // Set isDefinition=true to start searching right from the class level. - final PyClassTypeImpl classType = new PyClassTypeImpl(pyClass, true); - final List classAttrs = classType.resolveMember(name, target, AccessDirection.READ, resolveContext, true); - if (classAttrs == null) { - return null; - } - return StreamEx.of(classAttrs) - .map(RatedResolveResult::getElement) + } + else { + final Scope scope = ControlFlowCache.getScope(scopeOwner); + return StreamEx.of(scope.getNamedElements(name, false)) .select(PyTargetExpression.class) - .filter(x -> ScopeUtil.getScopeOwner(x) instanceof PyClass) .map(x -> getTypeFromTargetExpressionAnnotation(x, context)) .nonNull() .map(Ref::get) - .foldLeft(PyUnionType::union) + .findFirst() .orElse(null); } } diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java index c1a0a443192b..acc84fe4d304 100644 --- a/python/testSrc/com/jetbrains/python/PyTypingTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java @@ -720,6 +720,40 @@ public class PyTypingTest extends PyTestCase { " expr = x"); } + // PY-21864 + public void testLocalVariableAnnotationAheadOfTimeWithTarget() { + doTest("int", + "x: int\n" + + "with foo() as x:\n" + + " expr = x\n"); + } + + // PY-21864 + public void testLocalVariableAnnotationAheadOfTimeForTarget() { + doTest("int", + "x: int\n" + + "for x in foo():\n" + + " expr = x\n"); + } + + // PY-21864 + public void testLocalVariableAnnotationAheadOfTimeUnpackingTarget() { + doTest("int", + "x: int\n" + + "x, y = foo()\n" + + "expr = x"); + } + + // PY-21864 + public void testLocalVariableAnnotationAheadOfTimeOnlyFirstHintConsidered() { + doTest("int", + "x: int\n" + + "x = foo()\n" + + "x: str\n" + + "x = baz()\n" + + "expr = x"); + } + public void testInstanceAttributeAnnotation() { doTest("int", "class C:\n" +