From 46263763dcbd0a90d7dba5efc80d8f3fb5adf68f Mon Sep 17 00:00:00 2001 From: Dmitry Jemerov Date: Thu, 11 Aug 2011 18:41:33 +0200 Subject: [PATCH] named tuple support (PY-1360) --- python/src/com/jetbrains/python/PyNames.java | 3 + .../codeInsight/stdlib/PyNamedTupleType.java | 117 ++++++++++++++++++ .../stdlib/PyStdlibTypeProvider.java | 13 ++ .../python/psi/impl/PyCallExpressionImpl.java | 7 +- .../python/psi/types/PyCallableType.java | 19 +++ .../python/psi/types/PyClassType.java | 11 +- python/testData/completion/namedTuple.py | 5 + .../namedTuple.py | 6 + .../python/PythonCompletionTest.java | 9 ++ .../PyUnresolvedReferencesInspectionTest.java | 4 + 10 files changed, 187 insertions(+), 7 deletions(-) create mode 100644 python/src/com/jetbrains/python/codeInsight/stdlib/PyNamedTupleType.java create mode 100644 python/src/com/jetbrains/python/psi/types/PyCallableType.java create mode 100644 python/testData/completion/namedTuple.py create mode 100644 python/testData/inspections/PyUnresolvedReferencesInspection/namedTuple.py diff --git a/python/src/com/jetbrains/python/PyNames.java b/python/src/com/jetbrains/python/PyNames.java index f638bfd9def1..e625d353afa0 100644 --- a/python/src/com/jetbrains/python/PyNames.java +++ b/python/src/com/jetbrains/python/PyNames.java @@ -80,6 +80,9 @@ public class PyNames { public static final String NAME = "__name__"; public static final String ENTER = "__enter__"; + public static final String NAMEDTUPLE = "namedtuple"; + public static final String COLLECTIONS_PY = "collections.py"; + /** * Contains all known predefined names of "__foo__" form. */ diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyNamedTupleType.java b/python/src/com/jetbrains/python/codeInsight/stdlib/PyNamedTupleType.java new file mode 100644 index 000000000000..874c5802807f --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyNamedTupleType.java @@ -0,0 +1,117 @@ +package com.jetbrains.python.codeInsight.stdlib; + +import com.intellij.codeInsight.lookup.LookupElement; +import com.intellij.codeInsight.lookup.LookupElementBuilder; +import com.intellij.openapi.util.text.StringUtil; +import com.intellij.psi.PsiElement; +import com.intellij.util.ArrayUtil; +import com.intellij.util.ProcessingContext; +import com.jetbrains.python.psi.*; +import com.jetbrains.python.psi.impl.PyElementImpl; +import com.jetbrains.python.psi.resolve.PyResolveContext; +import com.jetbrains.python.psi.resolve.RatedResolveResult; +import com.jetbrains.python.psi.types.PyCallableType; +import com.jetbrains.python.psi.types.PyType; +import com.jetbrains.python.psi.types.TypeEvalContext; +import org.jetbrains.annotations.Nullable; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +/** + * @author yole + */ +public class PyNamedTupleType implements PyCallableType { + private final String myName; + private final boolean myDefinition; + private final PsiElement myDeclaration; + private final List myFields; + + public PyNamedTupleType(PsiElement declaration, String name, List fields, boolean isDefinition) { + myDeclaration = declaration; + myFields = fields; + myName = name; + myDefinition = isDefinition; + } + + @Override + public List resolveMember(String name, + @Nullable PyExpression location, + AccessDirection direction, + PyResolveContext resolveContext) { + if (hasField(name)) { + return Collections.singletonList(new RatedResolveResult(1000, new PyElementImpl(myDeclaration.getNode()))); + } + return Collections.emptyList(); + } + + private boolean hasField(String name) { + if (myDefinition) { + return "_make".equals(name); + } + else { + return myFields.contains(name) || "_replace".equals(name); + } + } + + @Override + public Object[] getCompletionVariants(String completionPrefix, PyExpression location, ProcessingContext context) { + List result = new ArrayList(); + if (!myDefinition) { + for (String field : myFields) { + result.add(LookupElementBuilder.create(field)); + } + } + return ArrayUtil.toObjectArray(result); + } + + @Override + public String getName() { + return "namedtuple '" + myName + "'"; + } + + @Override + public boolean isBuiltin(TypeEvalContext context) { + return false; + } + + @Override + public PyType getCallType() { + if (myDefinition) { + return new PyNamedTupleType(myDeclaration, myName, myFields, false); + } + return null; + } + + @Nullable + public static PyType fromCall(PyCallExpression call) { + final String name = PyUtil.strValue(call.getArgument(0, PyExpression.class)); + final PyExpression fieldNamesExpression = PyUtil.flattenParens(call.getArgument(1, PyExpression.class)); + if (name == null || fieldNamesExpression == null) { + return null; + } + List fieldNames = null; + if (fieldNamesExpression instanceof PySequenceExpression) { + fieldNames = PyUtil.strListValue(fieldNamesExpression); + } + else { + final String fieldNamesString = PyUtil.strValue(fieldNamesExpression); + if (fieldNamesString != null) { + fieldNames = parseFieldNamesString(fieldNamesString); + } + } + if (fieldNames != null) { + return new PyNamedTupleType(call, name, fieldNames, true); + } + return null; + } + + private static List parseFieldNamesString(String fieldNamesString) { + List result = new ArrayList(); + for(String name: StringUtil.tokenize(fieldNamesString, ", ")) { + result.add(name); + } + return result; + } +} diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java index 1350c6632cc9..41c502d29606 100644 --- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java @@ -4,6 +4,7 @@ import com.intellij.openapi.project.Project; import com.intellij.openapi.vfs.VirtualFile; import com.intellij.psi.PsiElement; import com.intellij.util.containers.HashMap; +import com.jetbrains.python.PyNames; import com.jetbrains.python.documentation.StructuredDocString; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyBuiltinCache; @@ -12,6 +13,7 @@ import com.jetbrains.python.psi.types.PyType; import com.jetbrains.python.psi.types.PyTypeParser; import com.jetbrains.python.psi.types.PyTypeProviderBase; import com.jetbrains.python.psi.types.TypeEvalContext; +import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; import java.io.IOException; @@ -28,6 +30,17 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { private Project myProject = null; private Map myTypeCache = new HashMap(); + @Override + public PyType getReferenceType(@NotNull PsiElement referenceTarget, TypeEvalContext context, @Nullable PsiElement anchor) { + if (referenceTarget instanceof PyFunction && + PyNames.NAMEDTUPLE.equals(((PyFunction) referenceTarget).getName()) && + PyNames.COLLECTIONS_PY.equals(referenceTarget.getContainingFile().getName()) && + anchor instanceof PyCallExpression) { + return PyNamedTupleType.fromCall((PyCallExpression) anchor); + } + return null; + } + @Override public PyType getReturnType(PyFunction function, @Nullable PyReferenceExpression callSite, TypeEvalContext context) { final String qname = getQualifiedName(function, callSite); diff --git a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java index 1de439ae4b72..218debf8a36e 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java @@ -119,11 +119,8 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress } else { final PyType type = context.getType(callee); - if (type instanceof PyClassType) { - PyClassType classType = (PyClassType)type; - if (classType.isDefinition()) { - return new PyClassType(classType.getPyClass(), false); - } + if (type instanceof PyCallableType) { + return ((PyCallableType) type).getCallType(); } return null; } diff --git a/python/src/com/jetbrains/python/psi/types/PyCallableType.java b/python/src/com/jetbrains/python/psi/types/PyCallableType.java new file mode 100644 index 000000000000..ee2fbdee4265 --- /dev/null +++ b/python/src/com/jetbrains/python/psi/types/PyCallableType.java @@ -0,0 +1,19 @@ +package com.jetbrains.python.psi.types; + +import org.jetbrains.annotations.Nullable; + +/** + * A type instances of which can possibly be called. For example, a class definition can be called, and the result of a call is a class + * instance. + * + * @author yole + */ +public interface PyCallableType extends PyType { + /** + * Returns the type which is the result of calling an instance of this type. + * + * @return the call result type or null if invalid. + */ + @Nullable + PyType getCallType(); +} diff --git a/python/src/com/jetbrains/python/psi/types/PyClassType.java b/python/src/com/jetbrains/python/psi/types/PyClassType.java index 96a96d6444ad..fd6b8e44fdca 100644 --- a/python/src/com/jetbrains/python/psi/types/PyClassType.java +++ b/python/src/com/jetbrains/python/psi/types/PyClassType.java @@ -10,7 +10,6 @@ import com.intellij.openapi.util.Pair; import com.intellij.openapi.util.UserDataHolderBase; import com.intellij.psi.PsiElement; import com.intellij.psi.util.PsiTreeUtil; -import com.intellij.psi.util.PsiUtilBase; import com.intellij.util.ArrayUtil; import com.intellij.util.ProcessingContext; import com.intellij.util.containers.ContainerUtil; @@ -31,7 +30,7 @@ import java.util.*; /** * @author yole */ -public class PyClassType extends UserDataHolderBase implements PyType { +public class PyClassType extends UserDataHolderBase implements PyCallableType { protected final PyClass myClass; protected final boolean myIsDefinition; @@ -188,6 +187,14 @@ public class PyClassType extends UserDataHolderBase implements PyType { return Collections.emptyList(); } + @Override + public PyType getCallType() { + if (isDefinition()) { + return new PyClassType(getPyClass(), false); + } + return null; + } + @Nullable private static PsiElement resolveClassMember(PyClassType aClass, String name, @Nullable PyExpression location) { PsiElement result = resolveInner(aClass.getPyClass(), name, location); diff --git a/python/testData/completion/namedTuple.py b/python/testData/completion/namedTuple.py new file mode 100644 index 000000000000..41efc4275b9d --- /dev/null +++ b/python/testData/completion/namedTuple.py @@ -0,0 +1,5 @@ +from collections import namedtuple + +Coord = namedtuple('Coord', 'lat long') +c = Coord(10, 20) +c. diff --git a/python/testData/inspections/PyUnresolvedReferencesInspection/namedTuple.py b/python/testData/inspections/PyUnresolvedReferencesInspection/namedTuple.py new file mode 100644 index 000000000000..097c0a60290f --- /dev/null +++ b/python/testData/inspections/PyUnresolvedReferencesInspection/namedTuple.py @@ -0,0 +1,6 @@ +from collections import namedtuple + +Point = namedtuple('Point', ['x', 'y'], verbose=True) + +p = Point(11, y=22) +print p.x + p.y diff --git a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java index cb324f96d7da..b988414eeab8 100644 --- a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java +++ b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java @@ -405,4 +405,13 @@ public class PythonCompletionTest extends PyLightFixtureTestCase { public void testMro() { // PY-3989 doTest(); } + + public void testNamedTuple() { // + final String testName = "completion/" + getTestName(true); + myFixture.configureByFile(testName + ".py"); + myFixture.completeBasic(); + final List strings = myFixture.getLookupElementStrings(); + assertTrue(strings.contains("lat")); + assertTrue(strings.contains("long")); + } } diff --git a/python/testSrc/com/jetbrains/python/inspections/PyUnresolvedReferencesInspectionTest.java b/python/testSrc/com/jetbrains/python/inspections/PyUnresolvedReferencesInspectionTest.java index 379a830ecf8b..8b53cffe9bf5 100644 --- a/python/testSrc/com/jetbrains/python/inspections/PyUnresolvedReferencesInspectionTest.java +++ b/python/testSrc/com/jetbrains/python/inspections/PyUnresolvedReferencesInspectionTest.java @@ -55,6 +55,10 @@ public class PyUnresolvedReferencesInspectionTest extends PyLightFixtureTestCase public void testBinaryOperators() { doTest(); } + + public void testNamedTuple() { + doTest(); + } private void doTest() { myFixture.configureByFile(TEST_DIRECTORY + getTestName(true) + ".py");