[ty] diagnostic on overridden __setattr__ and __delattr__ in frozen dataclasses

astral-sh/ty#111

Author: thejchap

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

@@ -452,8 +452,7 @@ from dataclasses import dataclass
 class MyFrozenClass:
     x: int
 
-    # TODO: Emit a diagnostic here
-    def __setattr__(self, name: str, value: object) -> None: ...
+    def __setattr__(self, name: str, value: object) -> None: ...  # error: [cannot-overwrite-attribute]
 
     # TODO: Emit a diagnostic here
     def __delattr__(self, name: str) -> None: ...

crates/ty_python_semantic/src/types/class.rs

@@ -19,7 +19,7 @@ use crate::semantic_index::{
 use crate::types::bound_super::BoundSuperError;
 use crate::types::constraints::{ConstraintSet, IteratorConstraintsExtension};
 use crate::types::context::InferContext;
-use crate::types::diagnostic::INVALID_TYPE_ALIAS_TYPE;
+use crate::types::diagnostic::{CANNOT_OVERWRITE_ATTRIBUTE, INVALID_TYPE_ALIAS_TYPE};
 use crate::types::enums::enum_metadata;
 use crate::types::function::{DataclassTransformerParams, KnownFunction};
 use crate::types::generics::{
@@ -2285,9 +2285,8 @@ impl<'db> ClassLiteral<'db> {
         inherited_generic_context: Option<GenericContext<'db>>,
         name: &str,
     ) -> Option<Type<'db>> {
-        let dataclass_params = self.dataclass_params(db);
-
         let field_policy = CodeGeneratorKind::from_class(db, self, specialization)?;
+        let dataclass_params = self.dataclass_params(db);
 
         let mut transformer_params =
             if let CodeGeneratorKind::DataclassLike(Some(transformer_params)) = field_policy {
@@ -2989,6 +2988,85 @@ impl<'db> ClassLiteral<'db> {
             .collect()
     }
 
+    pub(crate) fn validate_members(self, context: &InferContext<'db, '_>) {
+        let db = context.db();
+        let Some(field_policy) = CodeGeneratorKind::from_class(db, self, None) else {
+            return;
+        };
+        let dataclass_params = self.dataclass_params(db);
+        let mut transformer_params =
+            if let CodeGeneratorKind::DataclassLike(Some(transformer_params)) = field_policy {
+                Some(DataclassParams::from_transformer_params(
+                    db,
+                    transformer_params,
+                ))
+            } else {
+                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)
+                            {
+                                for (flag_name, flag) in DATACLASS_FLAGS {
+                                    if arg_name.as_str() == *flag_name {
+                                        flags.set(*flag, 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))
+        };
+        let db = context.db();
+        let class_body_scope = self.body_scope(db);
+        let table = place_table(db, class_body_scope);
+        let use_def = use_def_map(db, class_body_scope);
+        for (symbol_id, declarations) in use_def.all_end_of_scope_symbol_declarations() {
+            let result = place_from_declarations(db, declarations.clone());
+            let attr = result.ignore_conflicting_declarations();
+            let symbol = table.symbol(symbol_id);
+            let name = symbol.name();
+            if let Some(Type::FunctionLiteral(literal)) = attr.place.ignore_possibly_undefined()
+                && name == "__setattr__"
+            {
+                if let Some(CodeGeneratorKind::DataclassLike(_)) =
+                    CodeGeneratorKind::from_class(db, self, None)
+                    && has_dataclass_param(DataclassFlags::FROZEN)
+                {
+                    if let Some(builder) = context.report_lint(
+                        &CANNOT_OVERWRITE_ATTRIBUTE,
+                        literal.node(db, context.file(), context.module()),
+                    ) {
+                        let mut diagnostic = builder.into_diagnostic(format_args!(
+                            "Cannot overwrite attribute __setattr__ in class {}",
+                            self.name(db)
+                        ));
+                        diagnostic.info("__setattr__");
+                    }
+                }
+            }
+        }
+    }
+
     /// Returns a list of all annotated attributes defined in the body of this class. This is similar
     /// to the `__annotations__` attribute at runtime, but also contains default values.
     ///

crates/ty_python_semantic/src/types/diagnostic.rs

@@ -51,6 +51,7 @@ pub(crate) fn register_lints(registry: &mut LintRegistryBuilder) {
     registry.register_lint(&AMBIGUOUS_PROTOCOL_MEMBER);
     registry.register_lint(&CALL_NON_CALLABLE);
     registry.register_lint(&POSSIBLY_MISSING_IMPLICIT_CALL);
+    registry.register_lint(&CANNOT_OVERWRITE_ATTRIBUTE);
     registry.register_lint(&CONFLICTING_ARGUMENT_FORMS);
     registry.register_lint(&CONFLICTING_DECLARATIONS);
     registry.register_lint(&CONFLICTING_METACLASS);
@@ -392,6 +393,15 @@ declare_lint! {
     }
 }
 
+declare_lint! {
+    /// TODO
+    pub(crate) static CANNOT_OVERWRITE_ATTRIBUTE = {
+        summary: "TODO",
+        status: LintStatus::preview("1.0.0"),
+        default_level: Level::Error,
+    }
+}
+
 declare_lint! {
     /// ## What it does
     /// Checks for classes definitions which will fail at runtime due to

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

@@ -979,6 +979,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             if let Some(protocol) = class.into_protocol_class(self.db()) {
                 protocol.validate_members(&self.context);
             }
+
+            class.validate_members(&self.context);
         }
     }