named tuple support (PY-1360)

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.
    */

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<String> myFields;
+
+  public PyNamedTupleType(PsiElement declaration, String name, List<String> fields, boolean isDefinition) {
+    myDeclaration = declaration;
+    myFields = fields;
+    myName = name;
+    myDefinition = isDefinition;
+  }
+
+  @Override
+  public List<? extends RatedResolveResult> 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<LookupElement> result = new ArrayList<LookupElement>();
+    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<String> 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<String> parseFieldNamesString(String fieldNamesString) {
+    List<String> result = new ArrayList<String>();
+    for(String name: StringUtil.tokenize(fieldNamesString, ", ")) {
+      result.add(name);
+    }
+    return result;
+  }
+}

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<String, PyType> myTypeCache = new HashMap<String, PyType>();
 
+  @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);

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;
       }

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();
+}

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);

python/testData/completion/namedTuple.py

@@ -0,0 +1,5 @@
+from collections import namedtuple
+
+Coord = namedtuple('Coord', 'lat long')
+c = Coord(10, 20)
+c.<caret>

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

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<String> strings = myFixture.getLookupElementStrings();
+    assertTrue(strings.contains("lat"));
+    assertTrue(strings.contains("long"));
+  }
 }

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");