Override namedtuple._make and namedtuple._replace return type (PY-27148)

This commit is contained in:
Semyon Proshev
2018-01-17 22:06:35 +03:00
parent fcefe43db9
commit a81ded8e78
3 changed files with 124 additions and 18 deletions
@@ -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<PyClassLikeType> 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,
@@ -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)
}
}
@@ -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<TypeEvalContext> getTypeEvalContexts(@NotNull PyExpression element) {
return ImmutableList.of(TypeEvalContext.codeAnalysis(element.getProject(), element.getContainingFile()).withTracing(),
TypeEvalContext.userInitiated(element.getProject(), element.getContainingFile()).withTracing());