From a81ded8e782a098e591ea7403c9ac4fc171f2ca5 Mon Sep 17 00:00:00 2001 From: Semyon Proshev Date: Thu, 30 Nov 2017 19:26:17 +0300 Subject: [PATCH] Override `namedtuple._make` and `namedtuple._replace` return type (PY-27148) --- .../stdlib/PyStdlibTypeProvider.java | 21 ++++ .../stdlib/PyStdlibInspectionExtension.kt | 22 +---- .../com/jetbrains/python/PyTypeTest.java | 99 +++++++++++++++++++ 3 files changed, 124 insertions(+), 18 deletions(-) diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java index 7d04c185ba5b..d4bcc912e1b6 100644 --- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java @@ -252,6 +252,16 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { final PyClassLikeType classLikeType = as(firstArgument != null ? context.getType(firstArgument) : null, PyClassLikeType.class); return classLikeType != null ? Ref.create(classLikeType.toInstance()) : null; } + else if (callSite != null && + ArrayUtil.contains(qname, PyTypingTypeProvider.NAMEDTUPLE + "._make", PyTypingTypeProvider.NAMEDTUPLE + "._replace")) { + final PyExpression receiver = callSite.getReceiver(function); + if (receiver != null) { + final PyType receiverType = context.getType(receiver); + if (receiverType instanceof PyInstantiableType && isNamedTuple(receiverType, context)) { + return Ref.create(((PyInstantiableType)receiverType).toInstance()); + } + } + } } return null; @@ -336,6 +346,17 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { return null; } + public static boolean isNamedTuple(@Nullable PyType type, @NotNull TypeEvalContext context) { + if (type instanceof PyNamedTupleType) { + return true; + } + + final Condition isNT = + t -> t instanceof PyNamedTupleType || t != null && PyTypingTypeProvider.NAMEDTUPLE.equals(t.getClassQName()); + + return type instanceof PyClassLikeType && ContainerUtil.exists(((PyClassLikeType)type).getAncestorTypes(context), isNT); + } + @Nullable @Override public PyType getContextManagerVariableType(@NotNull PyClass contextManager, diff --git a/python/src/com/jetbrains/python/inspections/stdlib/PyStdlibInspectionExtension.kt b/python/src/com/jetbrains/python/inspections/stdlib/PyStdlibInspectionExtension.kt index a310283474ac..7a89f15cc49b 100644 --- a/python/src/com/jetbrains/python/inspections/stdlib/PyStdlibInspectionExtension.kt +++ b/python/src/com/jetbrains/python/inspections/stdlib/PyStdlibInspectionExtension.kt @@ -5,7 +5,7 @@ import com.intellij.psi.PsiReference import com.jetbrains.python.PyNames import com.jetbrains.python.codeInsight.stdlib.PyNamedTupleType import com.jetbrains.python.codeInsight.stdlib.PyStdlibClassMembersProvider -import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider +import com.jetbrains.python.codeInsight.stdlib.PyStdlibTypeProvider import com.jetbrains.python.inspections.PyInspectionExtension import com.jetbrains.python.psi.PyElement import com.jetbrains.python.psi.PyFunction @@ -47,22 +47,8 @@ class PyStdlibInspectionExtension : PyInspectionExtension() { override fun ignoreProtectedSymbol(expression: PyReferenceExpression, context: TypeEvalContext): Boolean { val qualifier = expression.qualifier - - if (qualifier != null && expression.referencedName in NAMEDTUPLE_SPECIAL_ATTRIBUTES) { - val qualifierType = context.getType(qualifier) - - if (qualifierType is PyNamedTupleType) { - return true - } - - val isTypingNT: (PyClassLikeType?) -> Boolean = - { it is PyNamedTupleType || it != null && PyTypingTypeProvider.NAMEDTUPLE == it.classQName } - - if (qualifierType is PyClassLikeType && qualifierType.getAncestorTypes(context).find(isTypingNT) != null) { - return true - } - } - - return false + return qualifier != null && + expression.referencedName in NAMEDTUPLE_SPECIAL_ATTRIBUTES && + PyStdlibTypeProvider.isNamedTuple(context.getType(qualifier), context) } } \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyTypeTest.java b/python/testSrc/com/jetbrains/python/PyTypeTest.java index b14e39960b9c..19bff2fee8cf 100644 --- a/python/testSrc/com/jetbrains/python/PyTypeTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypeTest.java @@ -2594,6 +2594,105 @@ public class PyTypeTest extends PyTestCase { "expr = compile(\"str\")"); } + // PY-27148 + public void testCollectionsNTMake() { + doTest("Cat", + "from collections import namedtuple\n" + + "Cat = namedtuple(\"Cat\", \"name age\")\n" + + "expr = Cat(\"name\", 5)._make([\"newname\", 6])"); + + doTest("Cat", + "from collections import namedtuple\n" + + "Cat = namedtuple(\"Cat\", \"name age\")\n" + + "expr = Cat._make([\"newname\", 6])"); + + doTest("Cat", + "from collections import namedtuple\n" + + "class Cat(namedtuple(\"Cat\", \"name age\")):\n" + + " pass\n" + + "expr = Cat(\"name\", 5)._make([\"newname\", 6])"); + + doTest("Cat", + "from collections import namedtuple\n" + + "class Cat(namedtuple(\"Cat\", \"name age\")):\n" + + " pass\n" + + "expr = Cat._make([\"newname\", 6])"); + } + + // PY-27148 + public void testTypingNTMake() { + runWithLanguageLevel( + LanguageLevel.PYTHON36, + () -> doTest("Cat", + "from typing import NamedTuple\n" + + "class Cat(NamedTuple):\n" + + " name: str\n" + + " age: int\n" + + "expr = Cat(\"name\", 5)._make([\"newname\", 6])") + ); + + runWithLanguageLevel( + LanguageLevel.PYTHON36, + () -> doTest("Cat", + "from typing import NamedTuple\n" + + "class Cat(NamedTuple):\n" + + " name: str\n" + + " age: int\n" + + "expr = Cat._make([\"newname\", 6])") + ); + + runWithLanguageLevel( + LanguageLevel.PYTHON36, + () -> doTest("Cat", + "from typing import NamedTuple\n" + + "Cat = NamedTuple(\"Cat\", name=str, age=int)\n" + + "expr = Cat(\"name\", 5)._make([\"newname\", 6])") + ); + + runWithLanguageLevel( + LanguageLevel.PYTHON36, + () -> doTest("Cat", + "from typing import NamedTuple\n" + + "Cat = NamedTuple(\"Cat\", name=str, age=int)\n" + + "expr = Cat._make([\"newname\", 6])") + ); + } + + // PY-27148 + public void testCollectionsNTReplace() { + doTest("Cat", + "from collections import namedtuple\n" + + "Cat = namedtuple(\"Cat\", \"name age\")\n" + + "expr = Cat(\"name\", 5)._replace(name=\"newname\")"); + + doTest("Cat", + "from collections import namedtuple\n" + + "class Cat(namedtuple(\"Cat\", \"name age\")):\n" + + " pass\n" + + "expr = Cat(\"name\", 5)._replace(name=\"newname\")"); + } + + // PY-27148 + public void testTypingNTReplace() { + runWithLanguageLevel( + LanguageLevel.PYTHON36, + () -> doTest("Cat", + "from typing import NamedTuple\n" + + "class Cat(NamedTuple):\n" + + " name: str\n" + + " age: int\n" + + "expr = Cat(\"name\", 5)._replace(name=\"newname\")") + ); + + runWithLanguageLevel( + LanguageLevel.PYTHON36, + () -> doTest("Cat", + "from typing import NamedTuple\n" + + "Cat = NamedTuple(\"Cat\", name=str, age=int)\n" + + "expr = Cat(\"name\", 5)._replace(name=\"newname\")") + ); + } + private static List getTypeEvalContexts(@NotNull PyExpression element) { return ImmutableList.of(TypeEvalContext.codeAnalysis(element.getProject(), element.getContainingFile()).withTracing(), TypeEvalContext.userInitiated(element.getProject(), element.getContainingFile()).withTracing());