[ty] Support class-arguments for dataclass transformers (#21457)

## Summary

Allow metaclass-based and baseclass-based dataclass-transformers to
overwrite the default behavior using class arguments:

```py
class Person(Model, order=True):
    # ...
```

## Conformance tests

Four new tests passing!

## Test Plan

New Markdown tests
This commit is contained in:
David Peter 2025-11-15 17:47:48 +01:00 committed by GitHub
parent 698231a47a
commit 29acc1e860
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 63 additions and 8 deletions

View file

@ -356,13 +356,17 @@ model < model # No error
### Overwriting of default parameters on the dataclass-like class
In the following examples, we show how a model can overwrite the default parameters set by the
`dataclass_transform` decorator. In particular, we change from `frozen=True` to `frozen=False`, and
from `order=False` (default) to `order=True`:
#### Using function-based transformers
```py
from typing import dataclass_transform
@dataclass_transform(frozen_default=True)
def default_frozen_model(*, frozen: bool = True): ...
def default_frozen_model(*, frozen: bool = True, order: bool = False): ...
@default_frozen_model()
class Frozen:
name: str
@ -370,12 +374,16 @@ class Frozen:
f = Frozen(name="test")
f.name = "new" # error: [invalid-assignment]
@default_frozen_model(frozen=False)
Frozen(name="A") < Frozen(name="B") # error: [unsupported-operator]
@default_frozen_model(frozen=False, order=True)
class Mutable:
name: str
m = Mutable(name="test")
m.name = "new" # No error
reveal_type(Mutable(name="A") < Mutable(name="B")) # revealed: bool
```
#### Using metaclass-based transformers
@ -392,6 +400,7 @@ class DefaultFrozenMeta(type):
namespace,
*,
frozen: bool = True,
order: bool = False,
): ...
class DefaultFrozenModel(metaclass=DefaultFrozenMeta): ...
@ -402,12 +411,17 @@ class Frozen(DefaultFrozenModel):
f = Frozen(name="test")
f.name = "new" # error: [invalid-assignment]
class Mutable(DefaultFrozenModel, frozen=False):
Frozen(name="A") < Frozen(name="B") # error: [unsupported-operator]
class Mutable(DefaultFrozenModel, frozen=False, order=True):
name: str
m = Mutable(name="test")
# TODO: no error here
# TODO: This should not be an error. In order to support this, we need to implement the precise `frozen` semantics of
# `dataclass_transform` described here: https://typing.python.org/en/latest/spec/dataclasses.html#dataclass-semantics
m.name = "new" # error: [invalid-assignment]
reveal_type(Mutable(name="A") < Mutable(name="B")) # revealed: bool
```
#### Using base-class-based transformers
@ -421,6 +435,7 @@ class DefaultFrozenModel:
cls,
*,
frozen: bool = True,
order: bool = False,
): ...
class Frozen(DefaultFrozenModel):
@ -429,12 +444,15 @@ class Frozen(DefaultFrozenModel):
f = Frozen(name="test")
f.name = "new" # error: [invalid-assignment]
class Mutable(DefaultFrozenModel, frozen=False):
Frozen(name="A") < Frozen(name="B") # error: [unsupported-operator]
class Mutable(DefaultFrozenModel, frozen=False, order=True):
name: str
m = Mutable(name="test")
# TODO: This should not be an error
m.name = "new" # error: [invalid-assignment]
m.name = "new" # No error
reveal_type(Mutable(name="A") < Mutable(name="B")) # revealed: bool
```
## `field_specifiers`

View file

@ -714,6 +714,13 @@ impl DefinitionKind<'_> {
}
}
pub(crate) const fn as_class(&self) -> Option<&AstNodeRef<ast::StmtClassDef>> {
match self {
DefinitionKind::Class(class) => Some(class),
_ => None,
}
}
pub(crate) fn is_import(&self) -> bool {
matches!(
self,

View file

@ -2206,7 +2206,7 @@ impl<'db> ClassLiteral<'db> {
let field_policy = CodeGeneratorKind::from_class(db, self, specialization)?;
let transformer_params =
let mut transformer_params =
if let CodeGeneratorKind::DataclassLike(Some(transformer_params)) = field_policy {
Some(DataclassParams::from_transformer_params(
db,
@ -2216,6 +2216,36 @@ impl<'db> ClassLiteral<'db> {
None
};
// Dataclass transformer flags can be overwritten using class arguments.
if let Some(transformer_params) = transformer_params.as_mut() {
if let Some(class_def) = self.definition(db).kind(db).as_class() {
let module = parsed_module(db, self.file(db)).load(db);
if let Some(arguments) = &class_def.node(&module).arguments {
let mut flags = transformer_params.flags(db);
for keyword in &arguments.keywords {
if let Some(arg_name) = &keyword.arg {
if let Some(is_set) =
keyword.value.as_boolean_literal_expr().map(|b| b.value)
{
match arg_name.as_str() {
"eq" => flags.set(DataclassFlags::EQ, is_set),
"order" => flags.set(DataclassFlags::ORDER, is_set),
"kw_only" => flags.set(DataclassFlags::KW_ONLY, is_set),
"frozen" => flags.set(DataclassFlags::FROZEN, is_set),
_ => {}
}
}
}
}
*transformer_params =
DataclassParams::new(db, flags, transformer_params.field_specifiers(db));
}
}
}
let has_dataclass_param = |param| {
dataclass_params.is_some_and(|params| params.flags(db).contains(param))
|| transformer_params.is_some_and(|params| params.flags(db).contains(param))