Support for subclasses of namedtuple defined in the current file (PY-4345)

python/psi-api/src/com/jetbrains/python/psi/PyClassRef.java

@@ -2,31 +2,41 @@ package com.jetbrains.python.psi;
 
 import com.intellij.psi.PsiElement;
 import com.jetbrains.python.psi.impl.PyQualifiedName;
+import com.jetbrains.python.psi.types.PyClassType;
 import org.jetbrains.annotations.Nullable;
 
 /**
  * @author yole
  */
 public class PyClassRef {
-  @Nullable
-  private final PsiElement myElement;
-
-  @Nullable
-  private final String myQName;
+  @Nullable private final PsiElement myElement;
+  @Nullable private final String myQName;
+  @Nullable private final PyClassType myType;
 
   public PyClassRef(@Nullable PsiElement element) {
     myElement = element;
     myQName = null;
+    myType = null;
   }
 
   public PyClassRef(@Nullable String qName) {
     myElement = null;
     myQName = qName;
+    myType = null;
+  }
+
+  public PyClassRef(@Nullable PyClassType type) {
+    myElement = null;
+    myQName = null;
+    myType = type;
   }
 
   @Nullable
   public PyClass getPyClass() {
-    return myElement instanceof PyClass ? (PyClass) myElement : null;
+    if (myElement instanceof PyClass) {
+      return (PyClass) myElement;
+    }
+    return null;
   }
 
   @Nullable
@@ -34,26 +44,40 @@ public class PyClassRef {
     return myElement;
   }
 
+  @Nullable
+  public PyClassType getType() {
+    return myType;
+  }
+
   @Nullable
   public String getClassName() {
     if (myElement instanceof PyClass) {
       return ((PyClass)myElement).getName();
     }
-    if (myQName != null) {
+    else if (myQName != null) {
       final PyQualifiedName qname = PyQualifiedName.fromDottedString(myQName);
       if (qname != null) {
         return qname.getLastComponent();
       }
     }
+    else if (myType != null) {
+      return myType.getName();
+    }
     return null;
   }
 
   @Nullable
   public String getQualifiedName() {
-    if (myQName != null) {
+    if (myElement instanceof PyClass) {
+      return ((PyClass)myElement).getQualifiedName();
+    }
+    else if (myQName != null) {
       return myQName;
     }
-    return myElement instanceof PyClass ? ((PyClass)myElement).getQualifiedName() : null;
+    else if (myType != null) {
+      return myType.getName();
+    }
+    return null;
   }
 
   @Override

python/src/com/jetbrains/python/codeInsight/stdlib/PyNamedTupleType.java

@@ -49,7 +49,7 @@ public class PyNamedTupleType extends PyClassTypeImpl implements PyCallableType
     if (myFields.contains(name)) {
       return Collections.singletonList(new RatedResolveResult(1000, new PyElementImpl(myDeclaration.getNode())));
     }
-    return Collections.emptyList();
+    return null;
   }
 
   @Override

python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java

@@ -60,6 +60,12 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase {
         }
       }
     }
+    else if (referenceTarget instanceof PyFunction && anchor instanceof PyCallExpression) {
+      if (PyNames.NAMEDTUPLE.equals(((PyFunction)referenceTarget).getName()) &&
+          PyNames.COLLECTIONS_PY.equals(referenceTarget.getContainingFile().getName())) {
+        return PyNamedTupleType.fromCall((PyCallExpression)anchor);
+      }
+    }
     return null;
   }
 

python/src/com/jetbrains/python/inspections/PyUnresolvedReferencesInspection.java

@@ -688,7 +688,7 @@ public class PyUnresolvedReferencesInspection extends PyInspection {
 
     private static boolean hasUnresolvedAncestors(PyClass cls) {
       for (PyClassRef classRef : cls.iterateAncestors()) {
-        if (classRef.getPyClass() == null) {
+        if (classRef.getPyClass() == null && classRef.getType() == null) {
           return true;
         }
       }

python/src/com/jetbrains/python/psi/impl/PyClassImpl.java

@@ -26,6 +26,7 @@ import com.jetbrains.python.psi.stubs.PropertyStubStorage;
 import com.jetbrains.python.psi.stubs.PyClassStub;
 import com.jetbrains.python.psi.stubs.PyFunctionStub;
 import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
+import com.jetbrains.python.psi.types.PyClassType;
 import com.jetbrains.python.psi.types.PyClassTypeImpl;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
@@ -148,20 +149,21 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
     final PyExpression[] superExpressions = getSuperClassExpressions();
     List<PsiElement> superClasses = new ArrayList<PsiElement>();
     for (PyExpression expr : superExpressions) {
-      expr = unfoldClass(expr);
-      if (expr instanceof PyReferenceExpression) {
-        final PsiPolyVariantReference ref = ((PyReferenceExpression)expr).getReference(PyResolveContext.noProperties());
-        final PsiElement result = ref.resolve();
-        if (result != null) {
-          superClasses.add(result);
-          continue;
-        }
-      }
-      superClasses.add(null);
+      superClasses.add(classElementFromExpression(expr));
     }
     return PsiUtilCore.toPsiElementArray(superClasses);
   }
 
+  @Nullable
+  public static PsiElement classElementFromExpression(@NotNull PyExpression expression) {
+    expression = unfoldClass(expression);
+    if (expression instanceof PyReferenceExpression) {
+      final PsiPolyVariantReference ref = ((PyReferenceExpression)expression).getReference(PyResolveContext.noProperties());
+      return ref.resolve();
+    }
+    return null;
+  }
+
   public static PyExpression unfoldClass(PyExpression expression) {
     if (expression instanceof PyCallExpression) {
       PyCallExpression call = (PyCallExpression)expression;
@@ -266,9 +268,19 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
     List<PyClassRef> result = resolveSuperClassesFromStub();
     if (result == null) {
       result = new ArrayList<PyClassRef>();
-      PsiElement[] superClassElements = getSuperClassElements();
-      for (PsiElement element : superClassElements) {
-        result.add(new PyClassRef(element));
+      final TypeEvalContext context = TypeEvalContext.fastStubOnly(null);
+      final PyExpression[] superClassExpressions = getSuperClassExpressions();
+      for (PyExpression expression : superClassExpressions) {
+        final PsiElement element = classElementFromExpression(expression);
+        if (element != null) {
+          result.add(new PyClassRef(element));
+        }
+        else {
+          final PyType type = expression.getType(context);
+          if (type instanceof PyClassType) {
+            result.add(new PyClassRef((PyClassType)type));
+          }
+        }
       }
     }
 

python/src/com/jetbrains/python/psi/types/PyClassTypeImpl.java

@@ -156,23 +156,25 @@ public class PyClassTypeImpl extends UserDataHolderBase implements PyClassType {
 
     for (PyClassRef superClass : myClass.iterateAncestors()) {
       final PyClass pyClass = superClass.getPyClass();
+      final PsiElement element = superClass.getElement();
+      final PyClassType type = superClass.getType();
       if (pyClass != null) {
         PsiElement superMember = resolveClassMember(pyClass, myIsDefinition, name, null);
         if (superMember != null) {
           return ResolveResultList.to(superMember);
         }
       }
-      else {
-        final PsiElement element = superClass.getElement();
-        if (element != null) {
-          for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
-            final PyType refType = typeProvider.getReferenceType(element, resolveContext.getTypeEvalContext(), myClass);
-            if (refType != null) {
-              return refType.resolveMember(name, location, direction, resolveContext);
-            }
+      else if (element != null) {
+        for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
+          final PyType refType = typeProvider.getReferenceType(element, resolveContext.getTypeEvalContext(), myClass);
+          if (refType != null) {
+            return refType.resolveMember(name, location, direction, resolveContext);
           }
         }
       }
+      else if (type != null) {
+        return type.resolveMember(name, location, direction, resolveContext);
+      }
     }
     if (isDefinition() && myClass.isNewStyleClass()) {
       PyClassType typeType = getMetaclassType();

python/testData/inspections/PyUnresolvedReferencesInspection/namedTuple.py

@@ -1,22 +0,0 @@
-from collections import namedtuple
-
-Point = namedtuple('Point', ['x', 'y'], verbose=True)
-
-print Point.x, Point.y
-
-p = Point(11, y=22)
-print p.x + p.y
-print p.__add__
-print p._asdict()
-print Point._fields
-print p._replace
-
-if isinstance(p, Point):
-    p.x
-
-class C(namedtuple('C', 'x')):
-    def f(self):
-        return self
-
-c = C()
-print(c.x, c.f())

python/testData/inspections/PyUnresolvedReferencesInspection3K/NamedTupleAssignment/a.py

@@ -0,0 +1,4 @@
+from b import Foo
+
+foo = Foo()
+print(foo.bar, foo.<warning descr="Unresolved attribute reference 'baz' for class 'Foo'">baz</warning>)

python/testData/inspections/PyUnresolvedReferencesInspection3K/NamedTupleAssignment/b.py

@@ -0,0 +1,4 @@
+from collections import namedtuple
+
+
+Foo = namedtuple('Foo', 'bar spam')

python/testData/inspections/PyUnresolvedReferencesInspection3K/NamedTupleBaseStub/a.py

@@ -0,0 +1,4 @@
+from b import Bar
+
+bar = Bar()
+print(bar.eggs, bar.<warning descr="Unresolved attribute reference 'bacon' for class 'Bar'">bacon</warning>)

python/testData/inspections/PyUnresolvedReferencesInspection3K/NamedTupleBaseStub/b.py

@@ -0,0 +1,5 @@
+from collections import namedtuple
+
+
+class Bar(namedtuple('Bar', 'spam eggs')):
+    pass

python/testData/inspections/PyUnresolvedReferencesInspection3K/namedTuple.py

@@ -0,0 +1,22 @@
+from collections import namedtuple
+
+Point = namedtuple('Point', ['x', 'y'], verbose=True)
+
+print(Point.x, Point.y)
+
+p = Point(11, y=22)
+print(p.x + p.y + p.<warning descr="Unresolved attribute reference 'z' for class 'Point'">z</warning>)
+print(p.__add__)
+print(p._asdict())
+print(Point._fields)
+print(p._replace)
+
+if isinstance(p, Point):
+    p.x
+
+class C(namedtuple('C', 'x y')):
+    def f(self):
+        return self
+
+c = C()
+print(c.x, c.y, c.<warning descr="Unresolved attribute reference 'z' for class 'C'">z</warning>, c.f())

python/testSrc/com/jetbrains/python/inspections/Py3UnresolvedReferencesInspectionTest.java

@@ -16,8 +16,19 @@ public class Py3UnresolvedReferencesInspectionTest extends PyTestCase {
     return ourPy3Descriptor;
   }
 
-  private void doMultiFileTest(@NotNull final String filename, @NotNull LanguageLevel level) {
-    runWithLanguageLevel(level, new Runnable() {
+  private void doTest() {
+    runWithLanguageLevel(LanguageLevel.PYTHON33, new Runnable() {
+      @Override
+      public void run() {
+        myFixture.configureByFile(TEST_DIRECTORY + getTestName(true) + ".py");
+        myFixture.enableInspections(PyUnresolvedReferencesInspection.class);
+        myFixture.checkHighlighting(true, false, false);
+      }
+    });
+  }
+
+  private void doMultiFileTest(@NotNull final String filename) {
+    runWithLanguageLevel(LanguageLevel.PYTHON33, new Runnable() {
       @Override
       public void run() {
         final String testName = getTestName(false);
@@ -29,11 +40,15 @@ public class Py3UnresolvedReferencesInspectionTest extends PyTestCase {
     });
   }
 
-  private void doMultiFileTest(@NotNull String filename) {
-    doMultiFileTest(filename, LanguageLevel.PYTHON33);
+  public void testNamedTuple() {
+    doTest();
   }
 
-  public void testNamedTupleStub() {
+  public void testNamedTupleAssignment() {
+    doMultiFileTest("a.py");
+  }
+
+  public void testNamedTupleBaseStub() {
     doMultiFileTest("a.py");
   }
 }

python/testSrc/com/jetbrains/python/inspections/PyUnresolvedReferencesInspectionTest.java

@@ -58,10 +58,6 @@ public class PyUnresolvedReferencesInspectionTest extends PyTestCase {
     doTest();
   }
 
-  public void testNamedTuple() {
-    doTest();
-  }
-
   // PY-2308
   public void testTypeAssertions() {
     doTest();