Support dataclasses fields defined via field function (PY-27398)

Update PyDataclassesTypeProvider to ignore such fields or correctly specify default value for them.
Create custom target stubs for such fields to store parameters.

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

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+ */
+package com.jetbrains.python.psi.stubs;
+
+import com.jetbrains.python.psi.impl.stubs.CustomTargetExpressionStub;
+
+public interface PyDataclassFieldStub extends CustomTargetExpressionStub {
+
+  /**
+   * @return true if `default` parameter is specified, false otherwise.
+   */
+  boolean hasDefault();
+
+  /**
+   * @return true if `default_factory` parameter is specified, false otherwise.
+   */
+  boolean hasDefaultFactory();
+
+  /**
+   * @return value of `init` parameter.
+   */
+  boolean initValue();
+}

python/src/META-INF/python-core-common.xml

@@ -700,6 +700,7 @@
     <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PropertyStubType"/>
     <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PyNamedTupleStubType"/>
     <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PyTypingAliasStubType"/>
+    <customTargetExpressionStubType implementation="com.jetbrains.python.psi.impl.stubs.PyDataclassFieldStubType"/>
 
     <dialectsTokenSetContributor implementation="com.jetbrains.python.PythonTokenSetContributor"/>
     <pyClassMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibClassMembersProvider"/>

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

@@ -7,7 +7,9 @@ import com.intellij.openapi.util.Ref
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
 import com.jetbrains.python.psi.*
 import com.jetbrains.python.psi.impl.PyCallExpressionNavigator
+import com.jetbrains.python.psi.impl.stubs.PyDataclassFieldStubImpl
 import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.stubs.PyDataclassFieldStub
 import com.jetbrains.python.psi.types.*
 
 class PyDataclassesTypeProvider : PyTypeProviderBase() {
@@ -64,13 +66,7 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() {
 
     cls.processClassLevelDeclarations { element, _ ->
       if (element is PyTargetExpression && !PyTypingTypeProvider.isClassVar(element, context)) {
-        val value = when {
-          context.maySwitchToAST(element) -> element.findAssignedValue()
-          element.hasAssignedValue() -> ellipsis
-          else -> null
-        }
-
-        parameters.add(PyCallableParameterImpl.nonPsi(element.name, getTypeForParameter(element, context), value))
+        fieldToParameter(element, ellipsis, context)?.also { parameters.add(it) }
       }
 
       true
@@ -79,6 +75,30 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() {
     return PyCallableTypeImpl(parameters, context.getType(cls))
   }
 
+  private fun fieldToParameter(field: PyTargetExpression,
+                               ellipsis: PyNoneLiteralExpression,
+                               context: TypeEvalContext): PyCallableParameter? {
+    val stub = field.stub
+    val fieldStub = if (stub == null) PyDataclassFieldStubImpl.create(field) else stub.getCustomStub(PyDataclassFieldStub::class.java)
+
+    return if (fieldStub == null) {
+      val value = when {
+        context.maySwitchToAST(field) -> field.findAssignedValue()
+        field.hasAssignedValue() -> ellipsis
+        else -> null
+      }
+
+      PyCallableParameterImpl.nonPsi(field.name, getTypeForParameter(field, context), value)
+    }
+    else if (!fieldStub.initValue()) {
+      null
+    }
+    else {
+      val value = if (fieldStub.hasDefault() || fieldStub.hasDefaultFactory()) ellipsis else null
+      PyCallableParameterImpl.nonPsi(field.name, getTypeForParameter(field, context), value)
+    }
+  }
+
   private fun getTypeForParameter(element: PyTargetExpression, context: TypeEvalContext): PyType? {
     val type = context.getType(element)
     if (type is PyCollectionType && type is PyClassType && type.classQName == DATACLASSES_INITVAR_TYPE) {

python/src/com/jetbrains/python/psi/PyFileElementType.java

@@ -62,7 +62,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
   @Override
   public int getStubVersion() {
     // Don't forget to update versions of indexes that use the updated stub-based elements
-    return 67;
+    return 68;
   }
 
   @Nullable

python/src/com/jetbrains/python/psi/impl/stubs/PyDataclassFieldStubImpl.kt

@@ -0,0 +1,108 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+ */
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.psi.stubs.StubInputStream
+import com.intellij.psi.stubs.StubOutputStream
+import com.intellij.psi.util.QualifiedName
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.impl.PyEvaluator
+import com.jetbrains.python.psi.resolve.PyResolveUtil
+import com.jetbrains.python.psi.stubs.PyDataclassFieldStub
+import java.io.IOException
+
+class PyDataclassFieldStubImpl private constructor(private val calleeName: QualifiedName,
+                                                   private val hasDefault: Boolean,
+                                                   private val hasDefaultFactory: Boolean,
+                                                   private val initValue: Boolean) : PyDataclassFieldStub {
+  companion object {
+    fun create(expression: PyTargetExpression): PyDataclassFieldStub? {
+      val value = expression.findAssignedValue() as? PyCallExpression ?: return null
+      val callee = value.callee as? PyReferenceExpression ?: return null
+
+      val calleeName = calculateFullyQCalleeName(callee) ?: calculateImportedCalleeName(callee) ?: return null
+      val arguments = analyzeArguments(value)
+
+      return PyDataclassFieldStubImpl(calleeName, arguments.first, arguments.second, arguments.third)
+    }
+
+    @Throws(IOException::class)
+    fun deserialize(stream: StubInputStream): PyDataclassFieldStub? {
+      val calleeName = stream.readName() ?: return null
+      val hasDefault = stream.readBoolean()
+      val hasDefaultFactory = stream.readBoolean()
+      val initValue = stream.readBoolean()
+
+      return PyDataclassFieldStubImpl(QualifiedName.fromDottedString(calleeName.string), hasDefault, hasDefaultFactory, initValue)
+    }
+
+    private fun calculateFullyQCalleeName(callee: PyReferenceExpression): QualifiedName? {
+      // SUPPORTED CASES:
+
+      // import dataclasses
+      // ... = dataclasses.field(...)
+
+      // import dataclasses as dc
+      // ... = dc.field(...)
+
+      val calleeName = callee.name
+      val qualifier = callee.qualifier
+
+      if (calleeName == "field" && qualifier is PyReferenceExpression && !qualifier.isQualified && resolvesToDataclassesModule(qualifier)) {
+        return QualifiedName.fromComponents(qualifier.name, calleeName)
+      }
+
+      return null
+    }
+
+    private fun calculateImportedCalleeName(callee: PyReferenceExpression): QualifiedName? {
+      // SUPPORTED CASES:
+
+      // from dataclasses import field
+      // ... = field(...)
+
+      // from dataclasses import field as F
+      // ... = F(...)
+
+      for (element in PyResolveUtil.resolveLocally(callee)) {
+        if (element is PyImportElement && element.importedQName.toString() == "field") {
+          val importStatement = element.containingImportStatement
+          if (importStatement is PyFromImportStatement && importStatement.importSourceQName.toString() == "dataclasses") {
+            return QualifiedName.fromComponents(callee.name)
+          }
+        }
+      }
+
+      return null
+    }
+
+    private fun analyzeArguments(call: PyCallExpression): Triple<Boolean, Boolean, Boolean> {
+      val hasDefault = call.getKeywordArgument("default") != null
+      val hasDefaultFactory = call.getKeywordArgument("default_factory") != null
+      val initValue = PyEvaluator().evaluate(call.getKeywordArgument("init")) as? Boolean ?: true
+
+      return Triple(hasDefault, hasDefaultFactory, initValue)
+    }
+
+    private fun resolvesToDataclassesModule(referenceExpression: PyReferenceExpression): Boolean {
+      return PyResolveUtil.resolveLocally(referenceExpression).any { it is PyImportElement && it.importedQName.toString() == "dataclasses" }
+    }
+  }
+
+  override fun getTypeClass(): Class<out CustomTargetExpressionStubType<out CustomTargetExpressionStub>> {
+    return PyDataclassFieldStubType::class.java
+  }
+
+  override fun serialize(stream: StubOutputStream) {
+    stream.writeName(calleeName.toString())
+    stream.writeBoolean(hasDefault)
+    stream.writeBoolean(hasDefaultFactory)
+    stream.writeBoolean(initValue)
+  }
+
+  override fun getCalleeName() = calleeName
+  override fun hasDefault() = hasDefault
+  override fun hasDefaultFactory() = hasDefaultFactory
+  override fun initValue() = initValue
+}

python/src/com/jetbrains/python/psi/impl/stubs/PyDataclassFieldStubType.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+ */
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.psi.stubs.StubInputStream
+import com.jetbrains.python.psi.PyTargetExpression
+import com.jetbrains.python.psi.stubs.PyDataclassFieldStub
+import java.io.IOException
+
+class PyDataclassFieldStubType : CustomTargetExpressionStubType<PyDataclassFieldStub>() {
+
+  override fun createStub(psi: PyTargetExpression): PyDataclassFieldStub? {
+    return PyDataclassFieldStubImpl.create(psi)
+  }
+
+  @Throws(IOException::class)
+  override fun deserializeStub(stream: StubInputStream): PyDataclassFieldStub? {
+    return PyDataclassFieldStubImpl.deserialize(stream)
+  }
+}

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

@@ -85,3 +85,18 @@ D1(<warning descr="Parameter 'a' unfilled"><warning descr="Parameter 'b' unfille
 D1(1<warning descr="Parameter 'b' unfilled">)</warning>
 D1(1, 2)
 D1(1, 2, <warning descr="Unexpected argument">3</warning>)
+
+
+@dataclasses.dataclass
+class E1:
+    a: int = dataclasses.field()
+    b: int = dataclasses.field(init=True)
+    c: int = dataclasses.field(init=False)
+    d: int = dataclasses.field(default=1)
+    e: int = dataclasses.field(default_factory=int)
+
+E1(1<warning descr="Parameter 'b' unfilled">)</warning>
+E1(1, 2)
+E1(1, 2, 3)
+E1(1, 2, 3, 4)
+E1(1, 2, 3, 4, <warning descr="Unexpected argument">5</warning>)

python/testData/inspections/PyArgumentListInspection/InitializingDataclass/dataclasses.py

@@ -1,3 +1,16 @@
+class _InitVarMeta(type):
+    def __getitem__(self, params):
+        return self
+
+class InitVar(metaclass=_InitVarMeta):
+    pass
+
+
 def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False,
               hash=None, frozen=False):
+    pass
+
+
+def field(*, default=_MISSING, default_factory=_MISSING, init=True, repr=True,
+          hash=None, compare=True, metadata=None):
     pass

python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/a.py

@@ -82,4 +82,22 @@ class D1:
     b: int
 
 D1(1, 2)
-D1(<warning descr="Expected type 'int', got 'str' instead">"1"</warning>, <warning descr="Expected type 'int', got 'str' instead">"2"</warning>)
+D1(<warning descr="Expected type 'int', got 'str' instead">"1"</warning>, <warning descr="Expected type 'int', got 'str' instead">"2"</warning>)
+
+
+@dataclasses.dataclass
+class E1:
+    a: int = dataclasses.field()
+    b: str = dataclasses.field(init=True)
+    c: int = dataclasses.field(init=False)
+    d: bytes = dataclasses.field(default=b"b")
+    e: int = dataclasses.field(default_factory=int)
+
+E1(1, "1")
+E1(<warning descr="Expected type 'int', got 'str' instead">"1"</warning>, <warning descr="Expected type 'str', got 'int' instead">1</warning>)
+
+E1(1, "1", b"1")
+E1(<warning descr="Expected type 'int', got 'bytes' instead">b"1"</warning>, "1", <warning descr="Expected type 'bytes', got 'int' instead">1</warning>)
+
+E1(1, "1", b"1", 1)
+E1(<warning descr="Expected type 'int', got 'str' instead">"1"</warning>, <warning descr="Expected type 'str', got 'bytes' instead">b"1"</warning>, <warning descr="Expected type 'bytes', got 'str' instead">"1"</warning>, <warning descr="Expected type 'int', got 'str' instead">"1"</warning>)

python/testData/inspections/PyTypeCheckerInspection/InitializingDataclass/dataclasses.py

@@ -8,4 +8,9 @@ class InitVar(metaclass=_InitVarMeta):
 
 def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False,
               hash=None, frozen=False):
+    pass
+
+
+def field(*, default=_MISSING, default_factory=_MISSING, init=True, repr=True,
+          hash=None, compare=True, metadata=None):
     pass

python/testData/paramInfo/InitializingDataclass/a.py

@@ -63,4 +63,15 @@ class D1:
     a: dataclasses.InitVar[int]
     b: int
 
-D1(<arg7>)
+D1(<arg7>)
+
+
+@dataclasses.dataclass
+class E1:
+    a: int = dataclasses.field()
+    b: int = dataclasses.field(init=True)
+    c: int = dataclasses.field(init=False)
+    d: int = dataclasses.field(default=1)
+    e: int = dataclasses.field(default_factory=int)
+
+E1(<arg8>)

python/testData/paramInfo/InitializingDataclass/dataclasses.py

@@ -8,4 +8,9 @@ class InitVar(metaclass=_InitVarMeta):
 
 def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False,
               hash=None, frozen=False):
+    pass
+
+
+def field(*, default=_MISSING, default_factory=_MISSING, init=True, repr=True,
+          hash=None, compare=True, metadata=None):
     pass

python/testData/stubs/dataclassField/a.py

@@ -0,0 +1,22 @@
+import dataclasses
+import dataclasses as dc
+from dataclasses import field
+from dataclasses import field as F
+from b import INIT_3
+
+
+INIT_0 = False
+INIT_1 = False
+INIT_2 = INIT_1
+
+
+@dataclasses.dataclass
+class A:
+    a: int = dataclasses.field(default=1)
+    b: int = dc.field(default_factory=int)
+    c: int = field(init=False)
+    d: int = F(init=True)
+    e: int = field(init=INIT_0)
+    f: int = field(init=INIT_2)
+    g: int = field(init=INIT_3)
+    h: int = field()

python/testData/stubs/dataclassField/b.py

@@ -0,0 +1 @@
+INIT_3 = False

python/testData/stubs/dataclassField/dataclasses.py

@@ -0,0 +1,8 @@
+def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False,
+              hash=None, frozen=False):
+    pass
+
+
+def field(*, default=_MISSING, default_factory=_MISSING, init=True, repr=True,
+          hash=None, compare=True, metadata=None):
+    pass

python/testSrc/com/jetbrains/python/PyParameterInfoTest.java

@@ -731,7 +731,7 @@ public class PyParameterInfoTest extends LightMarkedTestCase {
     runWithLanguageLevel(
       LanguageLevel.PYTHON37,
       () -> {
-        final Map<String, PsiElement> marks = loadMultiFileTest(7);
+        final Map<String, PsiElement> marks = loadMultiFileTest(8);
 
         feignCtrlP(marks.get("<arg1>").getTextOffset()).check("x: int, y: str, z: float=0.0", new String[]{"x: int, "});
         feignCtrlP(marks.get("<arg2>").getTextOffset()).check("x: int, y: str, z: float=0.0", new String[]{"x: int, "});
@@ -746,6 +746,8 @@ public class PyParameterInfoTest extends LightMarkedTestCase {
         feignCtrlP(marks.get("<arg5>").getTextOffset()).check("b: int", new String[]{"b: int"});
         feignCtrlP(marks.get("<arg6>").getTextOffset()).check("b: int", new String[]{"b: int"});
         feignCtrlP(marks.get("<arg7>").getTextOffset()).check("a: int, b: int", new String[]{"a: int, "});
+
+        feignCtrlP(marks.get("<arg8>").getTextOffset()).check("a: int, b: int, d: int=..., e: int=...", new String[]{"a: int, "});
       }
     );
   }

python/testSrc/com/jetbrains/python/PyStubsTest.java

@@ -409,10 +409,10 @@ public class PyStubsTest extends PyTestCase {
     final PyClass c = file.findTopLevelClass("C");
     assertNotNull(c);
     final TypeEvalContext context = TypeEvalContext.codeInsightFallback(myFixture.getProject());
-    assertNotNull(c.getMetaClassType(context));
+    assertNotNull(c.getMetaClassType(false, context));
     final PyClass d = file.findTopLevelClass("D");
     assertNotNull(d);
-    assertNotNull(d.getMetaClassType(context));
+    assertNotNull(d.getMetaClassType(false, context));
     assertNotParsed(file);
   }
 
@@ -924,4 +924,52 @@ public class PyStubsTest extends PyTestCase {
       assertNotNull(pyClass.findClassAttribute("foo", false, context));
     });
   }
+
+  // PY-27398
+  public void testDataclassField() {
+    class FieldChecker {
+
+      @NotNull
+      private final PyClass myClass;
+
+      private FieldChecker(@NotNull PyClass cls) {
+        myClass = cls;
+      }
+
+      private void check(@NotNull String name, boolean hasDefault, boolean hasDefaultFactory, boolean initValue) {
+        final TypeEvalContext context = TypeEvalContext.codeInsightFallback(myFixture.getProject());
+        final PyTargetExpression field = myClass.findClassAttribute(name, false, context);
+
+        final PyDataclassFieldStub fieldStub = field.getStub().getCustomStub(PyDataclassFieldStub.class);
+        assertNotNull(fieldStub);
+
+        assertEquals(hasDefault, fieldStub.hasDefault());
+        assertEquals(hasDefaultFactory, fieldStub.hasDefaultFactory());
+        assertEquals(initValue, fieldStub.initValue());
+      }
+    }
+
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON37,
+      () -> {
+        final PyFile file1 = getTestFile("dataclassField/a.py");
+        final PyFile file2 = getTestFile("dataclassField/dataclasses.py");
+        final PyFile file3 = getTestFile("dataclassField/b.py");
+
+        final FieldChecker checker = new FieldChecker(file1.findTopLevelClass("A"));
+        checker.check("a", true, false, true);
+        checker.check("b", false, true, true);
+        checker.check("c", false, false, false);
+        checker.check("d", false, false, true);
+        checker.check("e", false, false, false);
+        checker.check("f", false, false, false);
+        checker.check("g", false, false, true); // fallback `init` value
+        checker.check("h", false, false, true);
+
+        assertNotParsed(file1);
+        assertNotParsed(file2);
+        assertNotParsed(file3);
+      }
+    );
+  }
 }