diff --git a/crates/ty_python_semantic/resources/mdtest/dataclasses/dataclass_transform.md b/crates/ty_python_semantic/resources/mdtest/dataclasses/dataclass_transform.md index ebdf5073cd..8b7749359b 100644 --- a/crates/ty_python_semantic/resources/mdtest/dataclasses/dataclass_transform.md +++ b/crates/ty_python_semantic/resources/mdtest/dataclasses/dataclass_transform.md @@ -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` diff --git a/crates/ty_python_semantic/src/semantic_index/definition.rs b/crates/ty_python_semantic/src/semantic_index/definition.rs index 70a6039fd1..b01c112eec 100644 --- a/crates/ty_python_semantic/src/semantic_index/definition.rs +++ b/crates/ty_python_semantic/src/semantic_index/definition.rs @@ -714,6 +714,13 @@ impl DefinitionKind<'_> { } } + pub(crate) const fn as_class(&self) -> Option<&AstNodeRef> { + match self { + DefinitionKind::Class(class) => Some(class), + _ => None, + } + } + pub(crate) fn is_import(&self) -> bool { matches!( self, diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index 5b2295f066..ca83e6d981 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/crates/ty_python_semantic/src/types/class.rs @@ -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))