Reduce unstubbing in PyDataclassesTypeProvider and PyStdlibTypeProvider (PY-27398)

By introducing `hasAssignedValue` for target and its stub.

python/ipnb/src/org/jetbrains/plugins/ipnb/psi/IpnbPyTargetExpressionElementType.java

@@ -27,11 +27,13 @@ public class IpnbPyTargetExpressionElementType extends PyTargetExpressionElement
     super("IPNB_TARGET_EXPRESSION");
   }
 
+  @Override
   @NotNull
   public PsiElement createElement(@NotNull final ASTNode node) {
     return new IpnbPyTargetExpression(node);
   }
 
+  @Override
   public PyTargetExpression createPsi(@NotNull final PyTargetExpressionStub stub) {
     return new IpnbPyTargetExpression(stub);
   }

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

@@ -66,6 +66,18 @@ public interface PyTargetExpression extends PyQualifiedExpression, PsiNamedEleme
   @Nullable
   QualifiedName getCalleeName();
 
+  @Override
   @NotNull
   PsiReference getReference();
+
+  /**
+   * Checks if target has assigned value.
+   *
+   * This method does not access AST if underlying PSI is stub based.
+   *
+   * @return true if target has assigned expression, false otherwise (e.g. in type declaration statement).
+   */
+  default boolean hasAssignedValue() {
+    return true;
+  }
 }

python/psi-api/src/com/jetbrains/python/psi/stubs/PyTargetExpressionStub.java

@@ -63,4 +63,6 @@ public interface PyTargetExpressionStub extends NamedStub<PyTargetExpression>, P
 
   @Nullable
   String getDocString();
+
+  boolean hasAssignedValue();
 }

python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt

@@ -60,10 +60,17 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() {
     }
 
     val parameters = ArrayList<PyCallableParameter>()
+    val ellipsis = PyElementGenerator.getInstance(cls.project).createEllipsis()
 
     cls.processClassLevelDeclarations { element, _ ->
       if (element is PyTargetExpression && !PyTypingTypeProvider.isClassVar(element, context)) {
-        parameters.add(PyCallableParameterImpl.nonPsi(element.name, getTypeForParameter(element, context), element.findAssignedValue()))
+        val value = when {
+          context.maySwitchToAST(element) -> element.findAssignedValue()
+          element.hasAssignedValue() -> ellipsis
+          else -> null
+        }
+
+        parameters.add(PyCallableParameterImpl.nonPsi(element.name, getTypeForParameter(element, context), value))
       }
 
       true

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

@@ -630,9 +630,19 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase {
       }
     );
 
+    final PyNoneLiteralExpression ellipsis = PyElementGenerator.getInstance(cls.getProject()).createEllipsis();
+
     final Collector<PyTargetExpression, ?, LinkedHashMap<String, PyNamedTupleType.FieldTypeAndDefaultValue>> toNTFields =
       Collectors.toMap(PyTargetExpression::getName,
-                       field -> new PyNamedTupleType.FieldTypeAndDefaultValue(context.getType(field), field.findAssignedValue()),
+                       field -> {
+                         final PyExpression value = context.maySwitchToAST(field)
+                                                    ? field.findAssignedValue()
+                                                    : field.hasAssignedValue()
+                                                      ? ellipsis
+                                                      : null;
+
+                         return new PyNamedTupleType.FieldTypeAndDefaultValue(context.getType(field), value);
+                       },
                        (v1, v2) -> v2,
                        LinkedHashMap::new);
 

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

@@ -792,4 +792,13 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
   public String getTypeCommentAnnotation() {
     return getTypeCommentAnnotationFromStubOrPsi(this);
   }
+
+  @Override
+  public boolean hasAssignedValue() {
+    final PyTargetExpressionStub stub = getStub();
+    if (stub != null) {
+      return stub.hasAssignedValue();
+    }
+    return findAssignedValue() != null;
+  }
 }

python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java

@@ -32,7 +32,9 @@ import com.jetbrains.python.PythonDialectsTokenSetProvider;
 import com.jetbrains.python.documentation.docstrings.DocStringUtil;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyTargetExpressionImpl;
-import com.jetbrains.python.psi.stubs.*;
+import com.jetbrains.python.psi.stubs.PyFileStub;
+import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
+import com.jetbrains.python.psi.stubs.PyVariableNameIndex;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 
@@ -59,15 +61,18 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
     return myCustomStubTypes;
   }
 
+  @Override
   @NotNull
   public PsiElement createElement(@NotNull final ASTNode node) {
     return new PyTargetExpressionImpl(node);
   }
 
+  @Override
   public PyTargetExpression createPsi(@NotNull final PyTargetExpressionStub stub) {
     return new PyTargetExpressionImpl(stub);
   }
 
+  @Override
   @NotNull
   public PyTargetExpressionStub createStub(@NotNull final PyTargetExpression psi, final StubElement parentStub) {
     final String name = psi.getName();
@@ -79,7 +84,7 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
     for (CustomTargetExpressionStubType customStubType : getCustomStubTypes()) {
       CustomTargetExpressionStub customStub = customStubType.createStub(psi);
       if (customStub != null) {
-        return new PyTargetExpressionStubImpl(name, docString, typeComment, annotation, customStub, parentStub);
+        return new PyTargetExpressionStubImpl(name, docString, typeComment, annotation, psi.hasAssignedValue(), customStub, parentStub);
       }
     }
     PyTargetExpressionStub.InitializerType initializerType = PyTargetExpressionStub.InitializerType.Other;
@@ -95,9 +100,11 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
         initializer = ((PyReferenceExpression) callee).asQualifiedName();
       }
     }
-    return new PyTargetExpressionStubImpl(name, docString, initializerType, initializer, psi.isQualified(), typeComment, annotation, parentStub);
+    return new PyTargetExpressionStubImpl(name, docString, initializerType, initializer, psi.isQualified(), typeComment, annotation,
+                                          psi.hasAssignedValue(), parentStub);
   }
 
+  @Override
   public void serialize(@NotNull final PyTargetExpressionStub stub, @NotNull final StubOutputStream stream) throws IOException {
     stream.writeName(stub.getName());
     final String docString = stub.getDocString();
@@ -105,6 +112,7 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
     stream.writeVarInt(stub.getInitializerType().getIndex());
     stream.writeName(stub.getTypeComment());
     stream.writeName(stub.getAnnotation());
+    stream.writeBoolean(stub.hasAssignedValue());
     final CustomTargetExpressionStub customStub = stub.getCustomStub(CustomTargetExpressionStub.class);
     if (customStub != null) {
       stream.writeName(customStub.getTypeClass().getCanonicalName());
@@ -116,6 +124,7 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
     }
   }
 
+  @Override
   @NotNull
   public PyTargetExpressionStub deserialize(@NotNull final StubInputStream stream, final StubElement parentStub) throws IOException {
     String name = StringRef.toString(stream.readName());
@@ -128,21 +137,24 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
     final String typeComment = typeCommentRef == null ? null : typeCommentRef.getString();
     final StringRef annotationRef = stream.readName();
     final String annotation = annotationRef == null ? null : annotationRef.getString();
+    final boolean hasAssignedValue = stream.readBoolean();
     if (initializerType == PyTargetExpressionStub.InitializerType.Custom) {
       final String typeName = stream.readName().getString();
       for(CustomTargetExpressionStubType type: getCustomStubTypes()) {
         if (type.getClass().getCanonicalName().equals(typeName)) {
           CustomTargetExpressionStub stub = type.deserializeStub(stream);
-          return new PyTargetExpressionStubImpl(name, docString, typeComment, annotation, stub, parentStub);
+          return new PyTargetExpressionStubImpl(name, docString, typeComment, annotation, hasAssignedValue, stub, parentStub);
         }
       }
       throw new IOException("Unknown custom stub type " + typeName);
     }
     QualifiedName initializer = QualifiedName.deserialize(stream);
     boolean isQualified = stream.readBoolean();
-    return new PyTargetExpressionStubImpl(name, docString, initializerType, initializer, isQualified, typeComment, annotation, parentStub);
+    return new PyTargetExpressionStubImpl(name, docString, initializerType, initializer, isQualified, typeComment, annotation,
+                                          hasAssignedValue, parentStub);
   }
 
+  @Override
   public boolean shouldCreateStub(final ASTNode node) {
     if (PsiTreeUtil.getParentOfType(node.getPsi(), PyComprehensionElement.class, true, PyDocStringOwner.class) != null) {
       return false;

python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionStubImpl.java

@@ -33,6 +33,7 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
   private final boolean myQualified;
   private final String myTypeComment;
   private final String myAnnotation;
+  private final boolean myHasAssignedValue;
   
   @Nullable private final String myDocString;
   private final CustomTargetExpressionStub myCustomStub;
@@ -40,13 +41,15 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
   public PyTargetExpressionStubImpl(String name,
                                     @Nullable String docString,
                                     @Nullable String typeComment,
-                                    @Nullable String annotation, 
+                                    @Nullable String annotation,
+                                    boolean hasAssignedValue,
                                     CustomTargetExpressionStub customStub,
                                     StubElement parent) {
     super(parent, PyElementTypes.TARGET_EXPRESSION);
     myName = name;
     myTypeComment = typeComment;
     myAnnotation = annotation;
+    myHasAssignedValue = hasAssignedValue;
     myInitializerType = InitializerType.Custom;
     myInitializer = null;
     myQualified = false;
@@ -60,12 +63,14 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
                                     final QualifiedName initializer,
                                     final boolean qualified,
                                     @Nullable String typeComment, 
-                                    String annotation,
+                                    @Nullable String annotation,
+                                    boolean hasAssignedValue,
                                     final StubElement parentStub) {
     super(parentStub, PyElementTypes.TARGET_EXPRESSION);
     myName = name;
     myTypeComment = typeComment;
     myAnnotation = annotation;
+    myHasAssignedValue = hasAssignedValue;
     assert initializerType != InitializerType.Custom;
     myInitializerType = initializerType;
     myInitializer = initializer;
@@ -74,14 +79,17 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
     myDocString = docString;
   }
 
+  @Override
   public String getName() {
     return myName;
   }
 
+  @Override
   public InitializerType getInitializerType() {
     return myInitializerType;
   }
 
+  @Override
   public QualifiedName getInitializer() {
     return myInitializer;
   }
@@ -118,6 +126,11 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
     return myAnnotation;
   }
 
+  @Override
+  public boolean hasAssignedValue() {
+    return myHasAssignedValue;
+  }
+
   @Override
   public String toString() {
     return "PyTargetExpressionStub(name=" + myName + ")";

python/testData/inspections/PyArgumentListInspection/InitializingImportedTypingNamedTupleInheritor/a.py

@@ -0,0 +1,3 @@
+import b
+
+b.MyTup(5, "5")

python/testData/inspections/PyArgumentListInspection/InitializingImportedTypingNamedTupleInheritor/b.py

@@ -0,0 +1,6 @@
+import typing
+
+
+class MyTup(typing.NamedTuple):
+    bar: int
+    baz: str = ""

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

@@ -333,4 +333,8 @@ public class PyArgumentListInspectionTest extends PyInspectionTestCase {
   public void testInitializingDataclass() {
     runWithLanguageLevel(LanguageLevel.PYTHON37, this::doMultiFileTest);
   }
+
+  public void testInitializingImportedTypingNamedTupleInheritor() {
+    runWithLanguageLevel(LanguageLevel.PYTHON37, this::doMultiFileTest);
+  }
 }