[ty] Dataclass transform: neither frozen nor non-frozen

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

@@ -421,9 +421,7 @@ class Mutable(DefaultFrozenModel, frozen=False, order=True):
     name: str
 
 m = Mutable(name="test")
-# 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]
+m.name = "new"  # No error
 
 reveal_type(Mutable(name="A") < Mutable(name="B"))  # revealed: bool
 ```
@@ -459,6 +457,154 @@ m.name = "new"  # No error
 reveal_type(Mutable(name="A") < Mutable(name="B"))  # revealed: bool
 ```
 
+### Frozen inheritance
+
+Just like for regular `@dataclass`es, mixing frozen and non-frozen `@dataclass_transform` classes in
+an inheritance chain is not allowed. However, the root class of a `@dataclass_transform` hierarchy
+(the class decorated with `@dataclass_transform()` or the class that directly specifies the
+`@dataclass_transform` metaclass) is "neither frozen nor non-frozen", so both frozen and non-frozen
+subclasses can inherit from it.
+
+#### Using function-based transformers
+
+For function-based transformers, all classes are either frozen or non-frozen. There is no special
+root class.
+
+```py
+from typing import dataclass_transform
+
+@dataclass_transform(frozen_default=True)
+def frozen_model(*, frozen: bool = True): ...
+
+@frozen_model()
+class FrozenParent:
+    x: int
+
+@frozen_model()
+class FrozenChild(FrozenParent):
+    y: int
+
+@frozen_model(frozen=False)
+# error: [invalid-frozen-dataclass-subclass] "Non-frozen dataclass `NonFrozenChild` cannot inherit from frozen dataclass `FrozenParent`"
+class NonFrozenChild(FrozenParent):
+    y: int
+
+@frozen_model(frozen=False)
+class NonFrozenParent:
+    x: int
+
+@frozen_model()
+# error: [invalid-frozen-dataclass-subclass] "Frozen dataclass `FrozenFromNonFrozen` cannot inherit from non-frozen dataclass `NonFrozenParent`"
+class FrozenFromNonFrozen(NonFrozenParent):
+    y: int
+```
+
+#### Using metaclass-based transformers
+
+For metaclass-based transformers, the class that is decorated with `@dataclass_transform` is the
+root class that is "neither frozen nor non-frozen" (`DefaultFrozenMeta` in the example below). So
+children of that class can be either frozen or non-frozen:
+
+```py
+from typing import dataclass_transform
+
+@dataclass_transform(frozen_default=True)
+class FrozenMeta(type):
+    def __new__(
+        cls,
+        name,
+        bases,
+        namespace,
+        *,
+        frozen: bool = True,
+    ): ...
+
+class DefaultFrozenModel(metaclass=FrozenMeta): ...
+```
+
+Both frozen and non-frozen classes can inherit from the root class:
+
+```py
+class FrozenParent(DefaultFrozenModel):
+    x: int
+
+class NonFrozenParent(DefaultFrozenModel, frozen=False):
+    x: int
+```
+
+Inheriting from these classes is fine as long as we keep the frozen/non-frozen status consistent:
+
+```py
+class FrozenChild(FrozenParent):
+    y: int
+
+class NonFrozenChild(NonFrozenParent, frozen=False):
+    y: int
+```
+
+But mixing frozen and non-frozen is not allowed at this level:
+
+```py
+# error: [invalid-frozen-dataclass-subclass]
+class InvalidFrozenChild(NonFrozenParent, frozen=True):
+    y: int
+
+# error: [invalid-frozen-dataclass-subclass]
+class InvalidNonFrozenChild(FrozenParent, frozen=False):
+    y: int
+```
+
+#### Using base-class-based transformers
+
+Similarly, for base-class-based transformers, the class that is decorated with
+`@dataclass_transform` is the root class that is "neither frozen nor non-frozen"
+(`DefaultFrozenModel` in the example below). So children of that class can be either frozen or
+non-frozen:
+
+```py
+from typing import dataclass_transform
+
+@dataclass_transform(frozen_default=True)
+class DefaultFrozenModel:
+    def __init_subclass__(
+        cls,
+        *,
+        frozen: bool = True,
+    ): ...
+```
+
+Both frozen and non-frozen classes can inherit from that root model:
+
+```py
+class FrozenParent(DefaultFrozenModel):
+    x: int
+
+class NonFrozenParent(DefaultFrozenModel, frozen=False):
+    x: int
+```
+
+Inheriting from these classes is fine as long as we keep the frozen/non-frozen status consistent:
+
+```py
+class FrozenChild(FrozenParent):
+    y: int
+
+class NonFrozenChild(NonFrozenParent, frozen=False):
+    y: int
+```
+
+But mixing frozen and non-frozen is not allowed at this level:
+
+```py
+# error: [invalid-frozen-dataclass-subclass]
+class InvalidFrozenChild(NonFrozenParent, frozen=True):
+    y: int
+
+# error: [invalid-frozen-dataclass-subclass]
+class InvalidNonFrozenChild(FrozenParent, frozen=False):
+    y: int
+```
+
 ### Override diagnostics on dataclass-like classes
 
 #### Frozen override diagnostics
@@ -1206,7 +1352,7 @@ sure that we recognize all fields in a hierarchy like this:
 from dataclasses import dataclass
 from typing import dataclass_transform
 
-@dataclass_transform()
+@dataclass_transform(frozen_default=True)
 class ModelMeta(type):
     pass
 

crates/ty_python_semantic/src/types.rs

@@ -611,12 +611,6 @@ bitflags! {
     }
 }
 
-impl DataclassFlags {
-    pub(crate) const fn is_frozen(self) -> bool {
-        self.contains(Self::FROZEN)
-    }
-}
-
 pub(crate) const DATACLASS_FLAGS: &[(&str, DataclassFlags)] = &[
     ("init", DataclassFlags::INIT),
     ("repr", DataclassFlags::REPR),
@@ -12183,6 +12177,16 @@ pub(super) struct MetaclassCandidate<'db> {
     explicit_metaclass_of: StaticClassLiteral<'db>,
 }
 
+/// Information about a `@dataclass_transform`-decorated metaclass.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, salsa::Update, get_size2::GetSize)]
+pub(super) struct MetaclassTransformInfo<'db> {
+    pub(super) params: DataclassTransformerParams<'db>,
+
+    /// Whether the metaclass providing these parameters was declared on the class itself
+    /// (via an explicit `metaclass=` keyword) rather than inherited from a base class.
+    pub(super) from_explicit_metaclass: bool,
+}
+
 #[salsa::interned(debug, heap_size=ruff_memory_usage::heap_size)]
 pub struct UnionType<'db> {
     /// The union type includes values in any of these types.

crates/ty_python_semantic/src/types/class.rs

@@ -64,8 +64,8 @@ use crate::{
         semantic_index, use_def_map,
     },
     types::{
-        CallArguments, CallError, CallErrorKind, MetaclassCandidate, TypeDefinition, UnionType,
-        definition_expression_type,
+        CallArguments, CallError, CallErrorKind, MetaclassCandidate, MetaclassTransformInfo,
+        TypeDefinition, UnionType, definition_expression_type,
     },
 };
 use indexmap::IndexSet;
@@ -124,7 +124,7 @@ fn try_metaclass_cycle_initial<'db>(
     _db: &'db dyn Db,
     _id: salsa::Id,
     _self_: StaticClassLiteral<'db>,
-) -> Result<(Type<'db>, Option<DataclassTransformerParams<'db>>), MetaclassError<'db>> {
+) -> Result<(Type<'db>, Option<MetaclassTransformInfo<'db>>), MetaclassError<'db>> {
     Err(MetaclassError {
         kind: MetaclassErrorKind::Cycle,
     })
@@ -185,8 +185,8 @@ impl<'db> CodeGeneratorKind<'db> {
         ) -> Option<CodeGeneratorKind<'db>> {
             if class.dataclass_params(db).is_some() {
                 Some(CodeGeneratorKind::DataclassLike(None))
-            } else if let Ok((_, Some(transformer_params))) = class.try_metaclass(db) {
-                Some(CodeGeneratorKind::DataclassLike(Some(transformer_params)))
+            } else if let Ok((_, Some(info))) = class.try_metaclass(db) {
+                Some(CodeGeneratorKind::DataclassLike(Some(info.params)))
             } else if let Some(transformer_params) =
                 class.iter_mro(db, specialization).skip(1).find_map(|base| {
                     base.into_class().and_then(|class| {
@@ -2815,6 +2815,38 @@ impl<'db> StaticClassLiteral<'db> {
         (dataclass_params, transformer_params)
     }
 
+    /// Returns the effective frozen status of this class if it's a dataclass-like class.
+    ///
+    /// Returns `Some(true)` for a frozen dataclass-like class, `Some(false)` for a non-frozen one,
+    /// and `None` if the class is not a dataclass-like class, or if the dataclass is neither frozen
+    /// nor non-frozen.
+    pub(crate) fn is_frozen_dataclass(self, db: &'db dyn Db) -> Option<bool> {
+        // Check if this is a base-class-based transformer that has dataclass_transformer_params directly
+        // attached to it (because it is itself decorated with `@dataclass_transform`), or if this class
+        // has an explicit metaclass that is decorated with `@dataclass_transform`.
+        //
+        // In both cases, this signifies that this class is neither frozen nor non-frozen.
+        //
+        // See <https://typing.python.org/en/latest/spec/dataclasses.html#dataclass-semantics> for details.
+        if self.dataclass_transformer_params(db).is_some()
+            || self
+                .try_metaclass(db)
+                .is_ok_and(|(_, info)| info.is_some_and(|i| i.from_explicit_metaclass))
+        {
+            return None;
+        }
+
+        if let field_policy @ CodeGeneratorKind::DataclassLike(_) =
+            CodeGeneratorKind::from_class(db, self.into(), None)?
+        {
+            // Otherwise, if this class is a dataclass-like class, determine its frozen status based on
+            // dataclass params and dataclass transformer params.
+            Some(self.has_dataclass_param(db, field_policy, DataclassFlags::FROZEN))
+        } else {
+            None
+        }
+    }
+
     /// Checks if the given dataclass parameter flag is set for this class.
     /// This checks both the `dataclass_params` and `transformer_params`.
     fn has_dataclass_param(
@@ -2864,7 +2896,7 @@ impl<'db> StaticClassLiteral<'db> {
     pub(super) fn try_metaclass(
         self,
         db: &'db dyn Db,
-    ) -> Result<(Type<'db>, Option<DataclassTransformerParams<'db>>), MetaclassError<'db>> {
+    ) -> Result<(Type<'db>, Option<MetaclassTransformInfo<'db>>), MetaclassError<'db>> {
         tracing::trace!("StaticClassLiteral::try_metaclass: {}", self.name(db));
 
         // Identify the class's own metaclass (or take the first base class's metaclass).
@@ -2968,11 +3000,15 @@ impl<'db> StaticClassLiteral<'db> {
             });
         }
 
-        let dataclass_transformer_params = candidate
+        let transform_info = candidate
             .metaclass
             .static_class_literal(db)
-            .and_then(|(metaclass_literal, _)| metaclass_literal.dataclass_transformer_params(db));
-        Ok((candidate.metaclass.into(), dataclass_transformer_params))
+            .and_then(|(metaclass_literal, _)| metaclass_literal.dataclass_transformer_params(db))
+            .map(|params| MetaclassTransformInfo {
+                params,
+                from_explicit_metaclass: candidate.explicit_metaclass_of == self,
+            });
+        Ok((candidate.metaclass.into(), transform_info))
     }
 
     /// Returns the class member of this class named `name`.
@@ -3538,7 +3574,7 @@ impl<'db> StaticClassLiteral<'db> {
                 signature_from_fields(vec![self_parameter], instance_ty)
             }
             (CodeGeneratorKind::DataclassLike(_), "__setattr__") => {
-                if self.has_dataclass_param(db, field_policy, DataclassFlags::FROZEN) {
+                if self.is_frozen_dataclass(db) == Some(true) {
                     let signature = Signature::new(
                         Parameters::new(
                             db,
@@ -4033,7 +4069,7 @@ impl<'db> StaticClassLiteral<'db> {
             match name.as_str() {
                 "__setattr__" | "__delattr__" => {
                     if let CodeGeneratorKind::DataclassLike(_) = field_policy
-                        && self.has_dataclass_param(db, field_policy, DataclassFlags::FROZEN)
+                        && self.is_frozen_dataclass(db) == Some(true)
                     {
                         if let Some(builder) = context.report_lint(
                             &INVALID_DATACLASS_OVERRIDE,

crates/ty_python_semantic/src/types/diagnostic.rs

@@ -31,9 +31,7 @@ use crate::types::{
     ProtocolInstanceType, SpecialFormType, SubclassOfInner, Type, TypeContext, binding_type,
     protocol_class::ProtocolClass,
 };
-use crate::types::{
-    DataclassFlags, KnownInstanceType, MemberLookupPolicy, TypeVarInstance, UnionType,
-};
+use crate::types::{KnownInstanceType, MemberLookupPolicy, TypeVarInstance, UnionType};
 use crate::{Db, DisplaySettings, FxIndexMap, Program, declare_lint};
 use itertools::Itertools;
 use ruff_db::{
@@ -5566,7 +5564,7 @@ pub(super) fn report_bad_frozen_dataclass_inheritance<'db>(
     class_node: &ast::StmtClassDef,
     base_class: StaticClassLiteral<'db>,
     base_class_node: &ast::Expr,
-    base_class_params: DataclassFlags,
+    base_is_frozen: bool,
 ) {
     let db = context.db();
 
@@ -5576,7 +5574,7 @@ pub(super) fn report_bad_frozen_dataclass_inheritance<'db>(
         return;
     };
 
-    let mut diagnostic = if base_class_params.is_frozen() {
+    let mut diagnostic = if base_is_frozen {
         let mut diagnostic =
             builder.into_diagnostic("Non-frozen dataclass cannot inherit from frozen dataclass");
         diagnostic.set_concise_message(format_args!(

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -1061,24 +1061,20 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 }
 
                 if let Some((base_class_literal, _)) = base_class.static_class_literal(self.db())
-                    && let (Some(base_params), Some(class_params)) = (
-                        base_class_literal.dataclass_params(self.db()),
-                        class.dataclass_params(self.db()),
+                    && let (Some(base_is_frozen), Some(class_is_frozen)) = (
+                        base_class_literal.is_frozen_dataclass(self.db()),
+                        class.is_frozen_dataclass(self.db()),
                     )
+                    && base_is_frozen != class_is_frozen
                 {
-                    let base_params = base_params.flags(self.db());
-                    let class_is_frozen = class_params.flags(self.db()).is_frozen();
-
-                    if base_params.is_frozen() != class_is_frozen {
-                        report_bad_frozen_dataclass_inheritance(
-                            &self.context,
-                            class,
-                            class_node,
-                            base_class_literal,
-                            &class_node.bases()[i],
-                            base_params,
-                        );
-                    }
+                    report_bad_frozen_dataclass_inheritance(
+                        &self.context,
+                        class,
+                        class_node,
+                        base_class_literal,
+                        &class_node.bases()[i],
+                        base_is_frozen,
+                    );
                 }
             }