PY-54560 Fix handling of class-based field specifiers, add extra tests

GitOrigin-RevId: e411c00a4aeb12da59787c0273e678cab70b3e07

python/python-psi-impl/src/com/jetbrains/python/codeInsight/PyDataclasses.kt

@@ -682,26 +682,25 @@ fun resolveDataclassFieldParameters(
   }
   if (field.calleeName == null) return null
   val fieldSpecifierDeclaration = PyResolveUtil.resolveQualifiedNameInScope(field.calleeName!!, ScopeUtil.getScopeOwner(dataclass)!!, context)
-    .mapNotNull {
-      when (it) {
-        is PyFunction -> it
-        // TODO add a test for constructors
-        is PyClass -> it.findInitOrNew(true, context)
-        else -> null
-      }
-    }
+    .filterIsInstance<PyQualifiedNameOwner>()
     .firstOrNull { 
       val qualifiedName = it.qualifiedName
       qualifiedName != null && QualifiedName.fromDottedString(qualifiedName) in dataclassParams.fieldSpecifiers 
     }
   if (fieldSpecifierDeclaration == null) return null
+  val fieldSpecifierCallable = when (fieldSpecifierDeclaration) {
+    is PyClass -> fieldSpecifierDeclaration.findInitOrNew(true, context)
+    is PyFunction -> fieldSpecifierDeclaration
+    else -> null
+  }  
+  if (fieldSpecifierCallable == null) return null
   return PyDataclassFieldParameters(
     hasDefault = fieldStub?.hasDefault() ?: false,
     hasDefaultFactory = fieldStub?.hasDefaultFactory() ?: false,
     // TODO Should we delegate to dataclass parameters init here?
     // TODO support overloading init with Literal types
-    initValue = fieldStub?.initValue() ?: getArgumentDefault("init", fieldSpecifierDeclaration) ?: true,
-    kwOnly = fieldStub?.kwOnly() ?: getArgumentDefault("kw_only", fieldSpecifierDeclaration) ?: dataclassParams.kwOnly,
+    initValue = fieldStub?.initValue() ?: getArgumentDefault("init", fieldSpecifierCallable) ?: true,
+    kwOnly = fieldStub?.kwOnly() ?: getArgumentDefault("kw_only", fieldSpecifierCallable) ?: dataclassParams.kwOnly,
     alias = fieldStub?.alias,
   )
 }

python/testData/inspections/PyDataclassInspection/DataclassTransformFieldsOrder/a.py

@@ -1,7 +1,7 @@
 import dataclasses
 from typing import ClassVar
 
-from decorator import my_dataclass, my_field
+from decorator import my_dataclass, my_field, MyField
 
 
 @my_dataclass()
@@ -59,6 +59,18 @@ class B3:
     b: int = my_field()
 
 
+@my_dataclass()
+class B4:
+    <error descr="Fields with a default value must come after any fields without a default.">a</error>: int = MyField(default=1)
+    b: int = MyField()
+
+
+@my_dataclass()
+class B5:
+    <error descr="Fields with a default value must come after any fields without a default.">a</error>: int = MyField(default_factory=int)
+    b: int = MyField()
+
+
 @my_dataclass()
 class C1:
     x: int = dataclasses.MISSING
@@ -68,23 +80,26 @@ class C1:
 @my_dataclass()
 class C2:
     x: int = my_field(default=dataclasses.MISSING)
-    y: int
+    y: int = MyField(default=dataclasses.MISSING)
+    z: int
 
-C2(1, 2)
+C2(1, 2, 3)
 
 
 @my_dataclass()
 class C3:
     x: int = my_field(default_factory=dataclasses.MISSING)
-    y: int
+    y: int = MyField(default_factory=dataclasses.MISSING)
+    z: int
 
-C3(1, 2)
+C3(1, 2, 3)
 
 
 @my_dataclass()
 class D1:
     x: int = 0
     y: int = my_field(init=False)
+    z: int = MyField(init=False)
 
 
 @my_dataclass()

python/testData/inspections/PyDataclassInspection/DataclassTransformFieldsOrder/decorator.py

@@ -7,6 +7,11 @@ def my_field(**kwargs):
     ...
 
 
-@dataclass_transform(field_specifiers=(my_field,))
+class MyField:
+    def __init__(self, **kwargs):
+        ...
+
+
+@dataclass_transform(field_specifiers=(my_field, MyField,))
 def my_dataclass(**kwargs) -> Callable[[type[T]], type[T]]:
     ...

python/testData/inspections/PyDataclassInspection/DataclassTransformKwOnlyFieldOrderInInheritance/a.py

@@ -1,4 +1,5 @@
-from decorator import my_dataclass, my_dataclass_kw_only_default, my_field, my_field_kw_only_default
+from decorator import my_dataclass, my_dataclass_kw_only_default, my_field, my_field_kw_only_default, MyField, \
+    MyFieldKwOnlyDefault
 
 
 @my_dataclass(kw_only=True)
@@ -26,6 +27,7 @@ class C2(KwOnlyImplicitClassParam):
 @my_dataclass()
 class KwOnlyExplicitFieldParam:
     kw_only_default: int = my_field(default=42, kw_only=True)
+    kw_only_default2: int = MyField(default=42, kw_only=True)
     kw_only_no_default: int
 
 
@@ -37,6 +39,7 @@ class C3(KwOnlyExplicitFieldParam):
 @my_dataclass()
 class KwOnlyImplicitFieldParam:
     kw_only_default: int = my_field_kw_only_default(default=42)
+    kw_only_default2: int = MyFieldKwOnlyDefault(default=42)
     kw_only_no_default: int
 
 
@@ -68,8 +71,10 @@ class SubKwOnlyImplicitClassParam(BaseNotKwOnlyDefault):
 @my_dataclass()
 class SubKwOnlyExplicitFieldParam(BaseNotKwOnlyDefault):
     kw_only_no_default: int = my_field(kw_only=True)
+    kw_only_no_default2: int = MyField(kw_only=True)
 
 
 @my_dataclass()
 class SubKwOnlyExplicitFieldParam(BaseNotKwOnlyDefault):
     kw_only_no_default: int = my_field_kw_only_default()
+    kw_only_no_default2: int = MyFieldKwOnlyDefault()

python/testData/inspections/PyDataclassInspection/DataclassTransformKwOnlyFieldOrderInInheritance/decorator.py

@@ -8,15 +8,26 @@ def my_field(*, default=dataclasses.MISSING, kw_only=False):
     ...
 
 
+class MyField:
+    def __init__(self, default=dataclasses.MISSING, kw_only=False):
+        ...
+
+
 def my_field_kw_only_default(*, default=dataclasses.MISSING, kw_only=True):
     ...
 
 
-@dataclass_transform(field_specifiers=(my_field, my_field_kw_only_default))
+class MyFieldKwOnlyDefault:
+    def __init__(self, default=dataclasses.MISSING, kw_only=True):
+        ...
+
+
+@dataclass_transform(field_specifiers=(my_field, my_field_kw_only_default, MyField, MyFieldKwOnlyDefault))
 def my_dataclass(**kwargs) -> Callable[[type[T]], type[T]]:
     ...
 
 
-@dataclass_transform(kw_only_default=True, field_specifiers=(my_field, my_field_kw_only_default))
+@dataclass_transform(kw_only_default=True,
+                     field_specifiers=(my_field, my_field_kw_only_default, MyField, MyFieldKwOnlyDefault))
 def my_dataclass_kw_only_default(**kwargs) -> Callable[[type[T]], type[T]]:
     ...

python/testData/paramInfo/InitializingDataclassTransformFieldSpecifierKwOnlyArgumentBaseClassApiClassSpecifiers.py

@@ -0,0 +1,83 @@
+from typing import dataclass_transform
+
+
+class FieldKwOnlyDefaultFalse:
+    def __init__(self, kw_only: bool = False):
+        ...
+
+
+class FieldKwOnlyDefaultTrue:
+    def __init__(self, kw_only: bool = True):
+        ...
+
+
+class FieldNoOwnParams:
+    def __init__(self):
+        ...
+
+
+@dataclass_transform(field_specifiers=(FieldKwOnlyDefaultFalse, FieldKwOnlyDefaultTrue, FieldNoOwnParams))
+class MyDataclass:
+    def __init_subclass__(cls, **kwargs):
+        ...
+
+
+@dataclass_transform(kw_only_default=True,
+                     field_specifiers=(FieldKwOnlyDefaultFalse, FieldKwOnlyDefaultTrue, FieldNoOwnParams))
+class MyDataclassKwOnlyDefaultTrue:
+    def __init_subclass__(cls, **kwargs):
+        ...
+
+
+@dataclass_transform(kw_only_default=False,
+                     field_specifiers=(FieldKwOnlyDefaultFalse, FieldKwOnlyDefaultTrue, FieldNoOwnParams))
+class MyDataclassKwOnlyDefaultFalse:
+    def __init_subclass__(cls, **kwargs):
+        ...
+
+
+class DataclassKwOnlyTrue(MyDataclass, kw_only=True):
+    not_kw_only_spec_default: int = FieldKwOnlyDefaultFalse()
+    not_kw_only_spec_arg: int = FieldKwOnlyDefaultTrue(kw_only=False)
+    kw_only_inferred: int = FieldNoOwnParams()
+    kw_only_spec_default: int = FieldKwOnlyDefaultTrue()
+    kw_only_spec_arg: int = FieldKwOnlyDefaultFalse(kw_only=True)
+
+
+class DataclassKwOnlyDefaultTrue(MyDataclassKwOnlyDefaultTrue):
+    not_kw_only_spec_default: int = FieldKwOnlyDefaultFalse()
+    not_kw_only_spec_arg: int = FieldKwOnlyDefaultTrue(kw_only=False)
+    kw_only_inferred: int = FieldNoOwnParams()
+    kw_only_spec_default: int = FieldKwOnlyDefaultTrue()
+    kw_only_spec_arg: int = FieldKwOnlyDefaultFalse(kw_only=True)
+
+
+class DataclassKwOnlyFalse(MyDataclass, kw_only=False):
+    not_kw_only_spec_default: int = FieldKwOnlyDefaultFalse()
+    not_kw_only_spec_arg: int = FieldKwOnlyDefaultTrue(kw_only=False)
+    not_kw_only_inferred: int = FieldNoOwnParams()
+    kw_only_spec_default: int = FieldKwOnlyDefaultTrue()
+    kw_only_spec_arg: int = FieldKwOnlyDefaultFalse(kw_only=True)
+
+
+class DataclassKwOnlyDefaultFalse(MyDataclassKwOnlyDefaultFalse):
+    not_kw_only_spec_default: int = FieldKwOnlyDefaultFalse()
+    not_kw_only_spec_arg: int = FieldKwOnlyDefaultTrue(kw_only=False)
+    not_kw_only_inferred: int = FieldNoOwnParams()
+    kw_only_spec_default: int = FieldKwOnlyDefaultTrue()
+    kw_only_spec_arg: int = FieldKwOnlyDefaultFalse(kw_only=True)
+
+
+class DataclassImplicitKwOnlyFalse(MyDataclass):
+    not_kw_only_spec_default: int = FieldKwOnlyDefaultFalse()
+    not_kw_only_spec_arg: int = FieldKwOnlyDefaultTrue(kw_only=False)
+    not_kw_only_inferred: int = FieldNoOwnParams()
+    kw_only_spec_default: int = FieldKwOnlyDefaultTrue()
+    kw_only_spec_arg: int = FieldKwOnlyDefaultFalse(kw_only=True)
+
+
+DataclassKwOnlyTrue(<arg1>)
+DataclassKwOnlyDefaultTrue(<arg2>)
+DataclassKwOnlyFalse(<arg3>)
+DataclassKwOnlyDefaultFalse(<arg4>)
+DataclassImplicitKwOnlyFalse(<arg5>)

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

@@ -1229,7 +1229,18 @@ public class PyParameterInfoTest extends LightMarkedTestCase {
   }
 
   // PY-54560
-  public void testInitializingDataclassTransformFieldSpecifierKwOnlyArgument() {
+  public void testInitializingDataclassTransformFieldSpecifierKwOnlyArgumentDecoratorApiFunctionSpecifiers() {
+    final Map<String, PsiElement> marks = loadTest(5);
+
+    feignCtrlP(marks.get("<arg1>").getTextOffset()).check("not_kw_only_spec_default: int, not_kw_only_spec_arg: int, *, kw_only_inferred: int, kw_only_spec_default: int, kw_only_spec_arg: int", new String[]{"not_kw_only_spec_default: int, "});
+    feignCtrlP(marks.get("<arg2>").getTextOffset()).check("not_kw_only_spec_default: int, not_kw_only_spec_arg: int, *, kw_only_inferred: int, kw_only_spec_default: int, kw_only_spec_arg: int", new String[]{"not_kw_only_spec_default: int, "});
+    feignCtrlP(marks.get("<arg3>").getTextOffset()).check("not_kw_only_spec_default: int, not_kw_only_spec_arg: int, not_kw_only_inferred: int, *, kw_only_spec_default: int, kw_only_spec_arg: int", new String[]{"not_kw_only_spec_default: int, "});
+    feignCtrlP(marks.get("<arg4>").getTextOffset()).check("not_kw_only_spec_default: int, not_kw_only_spec_arg: int, not_kw_only_inferred: int, *, kw_only_spec_default: int, kw_only_spec_arg: int", new String[]{"not_kw_only_spec_default: int, "});
+    feignCtrlP(marks.get("<arg5>").getTextOffset()).check("not_kw_only_spec_default: int, not_kw_only_spec_arg: int, not_kw_only_inferred: int, *, kw_only_spec_default: int, kw_only_spec_arg: int", new String[]{"not_kw_only_spec_default: int, "});
+  }
+
+  // PY-54560
+  public void testInitializingDataclassTransformFieldSpecifierKwOnlyArgumentBaseClassApiClassSpecifiers() {
     final Map<String, PsiElement> marks = loadTest(5);
 
     feignCtrlP(marks.get("<arg1>").getTextOffset()).check("not_kw_only_spec_default: int, not_kw_only_spec_arg: int, *, kw_only_inferred: int, kw_only_spec_default: int, kw_only_spec_arg: int", new String[]{"not_kw_only_spec_default: int, "});