From ae58f4e52acf79fc8e95a7422fd20e434c2b5241 Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Thu, 23 Nov 2017 14:18:14 +0300 Subject: [PATCH] PY-24729 Initial support of instance attributes annotated on class level --- .../typing/PyTypingTypeProvider.java | 69 ++++++++++++++++--- .../InstanceAttributeAnnotation.py | 5 ++ .../InstanceAttributeAnnotation.pyi | 2 + .../other.py | 5 ++ .../com/jetbrains/python/PyTypingTest.java | 67 +++++++++++++++--- .../com/jetbrains/python/pyi/PyiTypeTest.java | 5 ++ 6 files changed, 135 insertions(+), 18 deletions(-) create mode 100644 python/testData/pyi/type/instanceAttributeAnnotation/InstanceAttributeAnnotation.py create mode 100644 python/testData/pyi/type/instanceAttributeAnnotation/InstanceAttributeAnnotation.pyi create mode 100644 python/testData/types/AnnotatedInstanceAttributeInOtherFile/other.py diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java index b1a57c40d9e3..f9ce4ff9d365 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java @@ -19,6 +19,7 @@ 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.dataflow.scope.ScopeUtil; import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyFunctionTypeAnnotation; import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyFunctionTypeAnnotationFile; import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyParameterTypeList; @@ -31,6 +32,7 @@ import com.jetbrains.python.psi.impl.stubs.PyTypingAliasStubType; import com.jetbrains.python.psi.resolve.PyResolveContext; import com.jetbrains.python.psi.resolve.PyResolveImportUtil; import com.jetbrains.python.psi.resolve.PyResolveUtil; +import com.jetbrains.python.psi.resolve.RatedResolveResult; import com.jetbrains.python.psi.stubs.PyClassStub; import com.jetbrains.python.psi.types.*; import one.util.streamex.StreamEx; @@ -325,20 +327,71 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { if (PROTOCOL.equals(target.getQualifiedName())) { return createTypingProtocolType(); } - final PyExpression annotation = getAnnotationValue(target, context); - if (annotation != null) { - return Ref.deref(getType(annotation, new Context(context))); + final Ref annotatedType = getTypeFromTargetExpressionAnnotation(target, context); + if (annotatedType != null) { + return annotatedType.get(); } - final String comment = target.getTypeCommentAnnotation(); - if (comment != null) { - final PyType type = Ref.deref(getVariableTypeCommentType(comment, referenceTarget, new Context(context))); + + final PyClass pyClass = target.getContainingClass(); + // an assignment inside a method + if (pyClass != null && target.isQualified() && ScopeUtil.getScopeOwner(target) instanceof PyFunction) { + final String name = target.getReferencedName(); + final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context); + + final StreamEx typeStream; + if (context.maySwitchToAST(target)) { + final PyType qualifierType = context.getType(target.getQualifier()); + if (qualifierType instanceof PyUnionType) { + typeStream = StreamEx.of(((PyUnionType)qualifierType).getMembers()); + } + else { + typeStream = StreamEx.of(qualifierType); + } + } + else if (PyUtil.isInstanceAttribute(target)) { + typeStream = StreamEx.of(new PyClassTypeImpl(pyClass, false)); + } + else { + return null; + } + return typeStream + .select(PyClassLikeType.class) + .map(PyClassLikeType::toClass) // force search on the class level right away + .flatMap(x -> { + final List resolved = x.resolveMember(name, target, AccessDirection.READ, resolveContext, true); + return resolved == null ? StreamEx.empty() : StreamEx.of(resolved); + }) + .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); + } + } + return null; + } + + @Nullable + private static Ref getTypeFromTargetExpressionAnnotation(@NotNull PyTargetExpression target, @NotNull TypeEvalContext context) { + final PyExpression annotation = getAnnotationValue(target, context); + if (annotation != null) { + return getType(annotation, new Context(context)); + } + final String comment = target.getTypeCommentAnnotation(); + if (comment != null) { + final Ref fromTypeComment = getVariableTypeCommentType(comment, target, new Context(context)); + if (fromTypeComment != null) { + final PyType type = Ref.deref(fromTypeComment); if (type instanceof PyTupleType) { final PyTupleExpression tupleExpr = PsiTreeUtil.getParentOfType(target, PyTupleExpression.class); if (tupleExpr != null) { - return PyTypeChecker.getTargetTypeFromTupleAssignment(target, tupleExpr, (PyTupleType)type); + return Ref.create(PyTypeChecker.getTargetTypeFromTupleAssignment(target, tupleExpr, (PyTupleType)type)); } } - return type; + return fromTypeComment; } } return null; diff --git a/python/testData/pyi/type/instanceAttributeAnnotation/InstanceAttributeAnnotation.py b/python/testData/pyi/type/instanceAttributeAnnotation/InstanceAttributeAnnotation.py new file mode 100644 index 000000000000..1e6ca1a24b78 --- /dev/null +++ b/python/testData/pyi/type/instanceAttributeAnnotation/InstanceAttributeAnnotation.py @@ -0,0 +1,5 @@ +class C: + def __init__(self): + self.attr = None + +C().attr \ No newline at end of file diff --git a/python/testData/pyi/type/instanceAttributeAnnotation/InstanceAttributeAnnotation.pyi b/python/testData/pyi/type/instanceAttributeAnnotation/InstanceAttributeAnnotation.pyi new file mode 100644 index 000000000000..ec4ed19258f9 --- /dev/null +++ b/python/testData/pyi/type/instanceAttributeAnnotation/InstanceAttributeAnnotation.pyi @@ -0,0 +1,2 @@ +class C: + attr: int diff --git a/python/testData/types/AnnotatedInstanceAttributeInOtherFile/other.py b/python/testData/types/AnnotatedInstanceAttributeInOtherFile/other.py new file mode 100644 index 000000000000..265ecadc6d05 --- /dev/null +++ b/python/testData/types/AnnotatedInstanceAttributeInOtherFile/other.py @@ -0,0 +1,5 @@ +class C: + attr: int + + def __init__(self): + self.attr = 'foo' \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java index 403a91256b4e..c1a0a443192b 100644 --- a/python/testSrc/com/jetbrains/python/PyTypingTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java @@ -312,7 +312,7 @@ public class PyTypingTest extends PyTestCase { // PY-19220 public void testMultiLineAssignmentComment() { - doTest("List[str]", + doTest("List[str]", "from typing import List\n" + "\n" + "expr = [\n" + @@ -486,7 +486,7 @@ public class PyTypingTest extends PyTestCase { // PY-18726 public void testFunctionTypeCommentCallableParameter() { - doTest("(bool, str) -> int", + doTest("(bool, str) -> int", "from typing import Callable\n" + "\n" + "def f(cb):\n" + @@ -536,12 +536,12 @@ public class PyTypingTest extends PyTestCase { // PY-18598 public void testFunctionTypeCommentEllipsisParameters() { - doTest("(x: Any, y: Any, z: Any) -> int", + doTest("(x: Any, y: Any, z: Any) -> int", "def f(x, y=42, z='foo'):\n" + " # type: (...) -> int \n" + " pass\n" + "\n" + - "expr = f"); + "expr = f"); } // PY-20421 @@ -558,7 +558,7 @@ public class PyTypingTest extends PyTestCase { // PY-18762 public void testHomogeneousTuple() { - doTest("Tuple[int, ...]", + doTest("Tuple[int, ...]", "from typing import Tuple\n" + "\n" + "def f(xs: Tuple[int, ...]):\n" + @@ -578,7 +578,7 @@ public class PyTypingTest extends PyTestCase { // PY-18762 public void testHomogeneousTupleUnpackingTarget() { - doTest("int", + doTest("int", "from typing import Tuple\n" + "\n" + "xs = unknown() # type: Tuple[int, ...]\n" + @@ -618,7 +618,7 @@ public class PyTypingTest extends PyTestCase { // PY-18877 public void testFunctionTypeCommentOnTheSameLine() { - doTest("(x: int, y: int) -> None", + doTest("(x: int, y: int) -> None", "def f(x,\n" + " y): # type: (int, int) -> None\n" + " pass\n" + @@ -729,7 +729,7 @@ public class PyTypingTest extends PyTestCase { } public void testIllegalAnnotationTargets() { - doTest("Tuple[Any, int, Any, Any]", + doTest("Tuple[Any, int, Any, Any]", "(w, _): Tuple[int, Any]\n" + "((x)): int\n" + "y: bool = z = undefined()\n" + @@ -929,7 +929,7 @@ public class PyTypingTest extends PyTestCase { // PY-23053 public void testListContainingClasses() { - doTest("Type[str]", + doTest("Type[str]", "xs = [str]\n" + "expr = xs.pop()"); } @@ -985,6 +985,39 @@ public class PyTypingTest extends PyTestCase { " expr = x\n"); } + // PY-24729 + public void testAnnotatedInstanceAttributeReferenceOutsideClass() { + doTest("int", + "class C:\n" + + " attr: int\n" + + "\n" + + " def __init__(self):\n" + + " self.attr = 'foo'\n" + + "\n" + + "expr = C().attr\n"); + } + + // PY-24729 + public void testAnnotatedInstanceAttributeReferenceInsideClass() { + doTest("int", + "class C:\n" + + " attr: int\n" + + "\n" + + " def __init__(self):\n" + + " self.attr = 'foo'\n" + + " \n" + + " def m(self):\n" + + " expr = self.attr\n"); + } + + // PY-24729 + public void testAnnotatedInstanceAttributeInOtherFile() { + doMultiFileStubAwareTest("int", + "from other import C\n" + + "\n" + + "expr = C().attr"); + } + private void doTestNoInjectedText(@NotNull String text) { myFixture.configureByText(PythonFileType.INSTANCE, text); final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myFixture.getProject()); @@ -1008,9 +1041,23 @@ public class PyTypingTest extends PyTestCase { myFixture.copyDirectoryToProject("typing", ""); myFixture.configureByText(PythonFileType.INSTANCE, text); final PyExpression expr = myFixture.findElementByText("expr", PyExpression.class); - final TypeEvalContext codeAnalysis = TypeEvalContext.codeAnalysis(expr.getProject(),expr.getContainingFile()); + final TypeEvalContext codeAnalysis = TypeEvalContext.codeAnalysis(expr.getProject(), expr.getContainingFile()); final TypeEvalContext userInitiated = TypeEvalContext.userInitiated(expr.getProject(), expr.getContainingFile()).withTracing(); assertType("Failed in code analysis context", expectedType, expr, codeAnalysis); assertType("Failed in user initiated context", expectedType, expr, userInitiated); } + + private void doMultiFileStubAwareTest(@NotNull final String expectedType, @NotNull final String text) { + myFixture.copyDirectoryToProject("types/" + getTestName(false), ""); + myFixture.copyDirectoryToProject("typing", ""); + myFixture.configureByText(PythonFileType.INSTANCE, text); + final PyExpression expr = myFixture.findElementByText("expr", PyExpression.class); + + final TypeEvalContext codeAnalysis = TypeEvalContext.codeAnalysis(expr.getProject(), expr.getContainingFile()); + assertType("Failed in code analysis context", expectedType, expr, codeAnalysis); + assertProjectFilesNotParsed(expr.getContainingFile()); + + final TypeEvalContext userInitiated = TypeEvalContext.userInitiated(expr.getProject(), expr.getContainingFile()).withTracing(); + assertType("Failed in user initiated context", expectedType, expr, userInitiated); + } } diff --git a/python/testSrc/com/jetbrains/python/pyi/PyiTypeTest.java b/python/testSrc/com/jetbrains/python/pyi/PyiTypeTest.java index eeafd607ad3f..385b14d026ec 100644 --- a/python/testSrc/com/jetbrains/python/pyi/PyiTypeTest.java +++ b/python/testSrc/com/jetbrains/python/pyi/PyiTypeTest.java @@ -146,4 +146,9 @@ public class PyiTypeTest extends PyTestCase { public void testComparisonOperatorOverloads() { doTest("int"); } + + // PY-24929 + public void testInstanceAttributeAnnotation() { + doTest("int"); + } }