[ty] support kw_only=True for dataclasses

astral-sh/ty#111

Author: thejchap

crates/ty_python_semantic/resources/mdtest/dataclasses/dataclass_transform.md

@@ -195,7 +195,40 @@ OrderTrueOverwritten(1) < OrderTrueOverwritten(2)
 
 ### `kw_only_default`
 
-To do
+When provided, sets the default value for the `kw_only` parameter of `field()`.
+
+```py
+from typing import dataclass_transform
+from dataclasses import field
+
+@dataclass_transform(kw_only_default=True)
+def create_model(*, init=True): ...
+@create_model()
+class A:
+    name: str = field(default="Voldemort")
+
+a = A()
+a = A(name="Harry")
+a = A("Harry")  # error: [too-many-positional-arguments]
+```
+
+TODO: This can be overridden by the call to the decorator function.
+
+```py
+from typing import dataclass_transform
+
+@dataclass_transform(kw_only_default=True)
+def create_model(*, kw_only: bool = True): ...
+@create_model(kw_only=False)
+class CustomerModel:
+    id: int
+    name: str
+
+# TODO: Should not emit errors
+# error: [missing-argument]
+# error: [too-many-positional-arguments]
+c = CustomerModel(1, "Harry")
+```
 
 ### `field_specifiers`
 

crates/ty_python_semantic/resources/mdtest/dataclasses/dataclasses.md

@@ -465,7 +465,70 @@ To do
 
 ### `kw_only`
 
-To do
+An error is emitted if a dataclass is defined with `kw_only=True` and positional arguments are
+passed to the constructor.
+
+```toml
+[environment]
+python-version = "3.10"
+```
+
+```py
+from dataclasses import dataclass
+
+@dataclass(kw_only=True)
+class A:
+    x: int
+    y: int
+
+# error: [missing-argument] "No arguments provided for required parameters `x`, `y`"
+# error: [too-many-positional-arguments] "Too many positional arguments: expected 0, got 2"
+a = A(1, 2)
+a = A(x=1, y=2)
+```
+
+The class-level parameter can be overridden per-field.
+
+```py
+from dataclasses import dataclass, field
+
+@dataclass(kw_only=True)
+class A:
+    a: str = field(kw_only=False)
+    b: int = 0
+
+A("hi")
+```
+
+If some fields are `kw_only`, they should appear after all positional fields in the `__init__`
+signature.
+
+```py
+@dataclass
+class A:
+    b: int = field(kw_only=True, default=3)
+    a: str
+
+A("hi")
+```
+
+### `kw_only` - Python < 3.10
+
+For Python < 3.10, `kw_only` is not supported.
+
+```toml
+[environment]
+python-version = "3.9"
+```
+
+```py
+from dataclasses import dataclass
+
+@dataclass(kw_only=True)  # TODO: Emit a diagnostic here
+class A:
+    x: int
+    y: int
+```
 
 ### `slots`
 

crates/ty_python_semantic/resources/mdtest/dataclasses/fields.md

@@ -63,13 +63,12 @@ class Person:
     age: int | None = field(default=None, kw_only=True)
     role: str = field(default="user", kw_only=True)
 
-# TODO: the `age` and `role` fields should be keyword-only
-# revealed: (self: Person, name: str, age: int | None = None, role: str = Literal["user"]) -> None
+# revealed: (self: Person, name: str, *, age: int | None = None, role: str = Literal["user"]) -> None
 reveal_type(Person.__init__)
 
 alice = Person(role="admin", name="Alice")
 
-# TODO: this should be an error
+# error: [too-many-positional-arguments] "Too many positional arguments: expected 1, got 2"
 bob = Person("Bob", 30)
 ```
 

crates/ty_python_semantic/src/types.rs

@@ -6707,6 +6707,9 @@ pub struct FieldInstance<'db> {
 
     /// Whether this field is part of the `__init__` signature, or not.
     pub init: bool,
+
+    /// Whether or not this field can only be passed as a keyword argument to `__init__`.
+    pub kw_only: Option<bool>,
 }
 
 // The Salsa heap is tracked separately.
@@ -6722,6 +6725,7 @@ impl<'db> FieldInstance<'db> {
             db,
             self.default_type(db).normalized_impl(db, visitor),
             self.init(db),
+            self.kw_only(db),
         )
     }
 }

crates/ty_python_semantic/src/types/call/bind.rs

@@ -12,6 +12,7 @@ use ruff_db::parsed::parsed_module;
 use smallvec::{SmallVec, smallvec, smallvec_inline};
 
 use super::{Argument, CallArguments, CallError, CallErrorKind, InferContext, Signature, Type};
+use crate::Program;
 use crate::db::Db;
 use crate::dunder_all::dunder_all_names;
 use crate::place::{Boundness, Place};
@@ -33,7 +34,7 @@ use crate::types::{
     WrapperDescriptorKind, enums, ide_support, todo_type,
 };
 use ruff_db::diagnostic::{Annotation, Diagnostic, SubDiagnostic, SubDiagnosticSeverity};
-use ruff_python_ast as ast;
+use ruff_python_ast::{self as ast, PythonVersion};
 
 /// Binding information for a possible union of callables. At a call site, the arguments must be
 /// compatible with _all_ of the types in the union for the call to be valid.
@@ -860,7 +861,11 @@ impl<'db> Bindings<'db> {
                                     params |= DataclassParams::MATCH_ARGS;
                                 }
                                 if to_bool(kw_only, false) {
-                                    params |= DataclassParams::KW_ONLY;
+                                    if Program::get(db).python_version(db) >= PythonVersion::PY310 {
+                                        params |= DataclassParams::KW_ONLY;
+                                    } else {
+                                        // TODO: emit diagnostic
+                                    }
                                 }
                                 if to_bool(slots, false) {
                                     params |= DataclassParams::SLOTS;
@@ -919,7 +924,8 @@ impl<'db> Bindings<'db> {
                         }
 
                         Some(KnownFunction::Field) => {
-                            if let [default, default_factory, init, ..] = overload.parameter_types()
+                            if let [default, default_factory, init, .., kw_only] =
+                                overload.parameter_types()
                             {
                                 let default_ty = match (default, default_factory) {
                                     (Some(default_ty), _) => *default_ty,
@@ -933,6 +939,14 @@ impl<'db> Bindings<'db> {
                                     .map(|init| !init.bool(db).is_always_false())
                                     .unwrap_or(true);
 
+                                let kw_only = if Program::get(db).python_version(db)
+                                    >= PythonVersion::PY310
+                                {
+                                    kw_only.map(|kw_only| !kw_only.bool(db).is_always_false())
+                                } else {
+                                    None
+                                };
+
                                 // `typeshed` pretends that `dataclasses.field()` returns the type of the
                                 // default value directly. At runtime, however, this function returns an
                                 // instance of `dataclasses.Field`. We also model it this way and return
@@ -942,7 +956,7 @@ impl<'db> Bindings<'db> {
                                 // to `T`. Otherwise, we would error on `name: str = field(default="")`.
                                 overload.set_return_type(Type::KnownInstance(
                                     KnownInstanceType::Field(FieldInstance::new(
-                                        db, default_ty, init,
+                                        db, default_ty, init, kw_only,
                                     )),
                                 ));
                             }

crates/ty_python_semantic/src/types/class.rs

@@ -1114,6 +1114,9 @@ pub(crate) struct Field<'db> {
 
     /// Whether or not this field should appear in the signature of `__init__`.
     pub(crate) init: bool,
+
+    /// Whether or not this field can only be passed as a keyword argument to `__init__`.
+    pub(crate) kw_only: Option<bool>,
 }
 
 /// Representation of a class definition statement in the AST: either a non-generic class, or a
@@ -1866,6 +1869,7 @@ impl<'db> ClassLiteral<'db> {
                     mut default_ty,
                     init_only: _,
                     init,
+                    kw_only,
                 },
             ) in self.fields(db, specialization, field_policy)
             {
@@ -1930,7 +1934,10 @@ impl<'db> ClassLiteral<'db> {
                     }
                 }
 
-                let mut parameter = if kw_only_field_seen || name == "__replace__" {
+                let mut parameter = if kw_only_field_seen
+                    || name == "__replace__"
+                    || kw_only.unwrap_or(has_dataclass_param(DataclassParams::KW_ONLY))
+                {
                     Parameter::keyword_only(field_name)
                 } else {
                     Parameter::positional_or_keyword(field_name)
@@ -1949,6 +1956,10 @@ impl<'db> ClassLiteral<'db> {
                 parameters.push(parameter);
             }
 
+            // In the event that we have a mix of keyword-only and positional parameters, we need to sort them
+            // so that the keyword-only parameters appear after positional parameters.
+            parameters.sort_by_key(Parameter::is_keyword_only);
+
             let mut signature = Signature::new(Parameters::new(parameters), return_ty);
             signature.inherited_generic_context = self.generic_context(db);
             Some(CallableType::function_like(db, signature))
@@ -2261,9 +2272,11 @@ impl<'db> ClassLiteral<'db> {
                         default_ty.map(|ty| ty.apply_optional_specialization(db, specialization));
 
                     let mut init = true;
+                    let mut kw_only = None;
                     if let Some(Type::KnownInstance(KnownInstanceType::Field(field))) = default_ty {
                         default_ty = Some(field.default_type(db));
                         init = field.init(db);
+                        kw_only = field.kw_only(db);
                     }
 
                     attributes.insert(
@@ -2273,6 +2286,7 @@ impl<'db> ClassLiteral<'db> {
                             default_ty,
                             init_only: attr.is_init_var(),
                             init,
+                            kw_only,
                         },
                     );
                 }