fix decorated class not transformed #1180

crates/pyrefly_config/src/args.rs

@@ -19,6 +19,7 @@ use pyrefly_util::absolutize::Absolutize as _;
 use pyrefly_util::arc_id::ArcId;
 use pyrefly_util::display;
 
+use crate::base::ClassDecoratorBehavior;
 use crate::base::InferReturnTypes;
 use crate::base::Preset;
 use crate::base::RecursionOverflowHandler;
@@ -327,6 +328,9 @@ pub struct ConfigOverrideArgs {
     /// are only eligible when --check-unannotated-defs is true.
     #[arg(long)]
     infer_return_types: Option<InferReturnTypes>,
+    /// How class decorators affect class types: ignore, ignore-untyped (default), or apply.
+    #[arg(long)]
+    class_decorator_behavior: Option<ClassDecoratorBehavior>,
     /// Whether Pyrefly will respect ignore statements for other tools, e.g. `# pyright: ignore`.
     /// Equivalent to passing the names of all tools to `--enabled-ignores`.
     #[arg(
@@ -434,6 +438,7 @@ impl ConfigOverrideArgs {
             untyped_def_behavior,
             check_unannotated_defs,
             infer_return_types,
+            class_decorator_behavior,
             permissive_ignores,
             enabled_ignores,
             type_ignore_unknown_tag_behavior,
@@ -462,6 +467,7 @@ impl ConfigOverrideArgs {
             || untyped_def_behavior.is_some()
             || check_unannotated_defs.is_some()
             || infer_return_types.is_some()
+            || class_decorator_behavior.is_some()
             || permissive_ignores.is_some()
             || enabled_ignores.is_some()
             || type_ignore_unknown_tag_behavior.is_some()
@@ -565,6 +571,9 @@ impl ConfigOverrideArgs {
         if let Some(x) = &self.infer_return_types {
             config.root.infer_return_types = Some(*x);
         }
+        if let Some(x) = &self.class_decorator_behavior {
+            config.root.class_decorator_behavior = Some(*x);
+        }
         match (self.permissive_ignores, &self.enabled_ignores) {
             // Special case: if the underlying config sets enabled-ignores and --permissive-ignores
             // is passed on the command-line, we overwrite enabled-ignores.
@@ -759,6 +768,33 @@ mod tests {
         );
     }
 
+    #[test]
+    fn class_decorator_behavior_cli_overrides_config() {
+        for (value, expected) in [
+            ("ignore", ClassDecoratorBehavior::Ignore),
+            ("ignore-untyped", ClassDecoratorBehavior::IgnoreUntyped),
+            ("apply", ClassDecoratorBehavior::Apply),
+        ] {
+            let args =
+                ConfigOverrideArgs::parse_from(["pyrefly", "--class-decorator-behavior", value]);
+            assert!(args.has_overrides());
+            let mut config = ConfigFile::default();
+            config.interpreters.skip_interpreter_query = true;
+            config.root.class_decorator_behavior = Some(ClassDecoratorBehavior::Apply);
+            let (config, errors) = args.override_config(config);
+            assert!(errors.is_empty());
+            assert_eq!(config.root.class_decorator_behavior, Some(expected));
+        }
+        assert!(
+            ConfigOverrideArgs::try_parse_from([
+                "pyrefly",
+                "--class-decorator-behavior",
+                "invalid"
+            ])
+            .is_err()
+        );
+    }
+
     #[test]
     fn cli_interpreter_selection_overrides_config_skip() {
         for command_line in [

crates/pyrefly_config/src/base.rs

@@ -50,6 +50,29 @@ pub enum InferReturnTypes {
     Checked,
 }
 
+/// Controls how class decorators affect the type bound to the class name.
+#[derive(
+    Debug,
+    PartialEq,
+    Eq,
+    Deserialize,
+    Serialize,
+    Clone,
+    Copy,
+    Default,
+    ValueEnum
+)]
+#[serde(rename_all = "kebab-case")]
+pub enum ClassDecoratorBehavior {
+    /// Preserve the class regardless of decorator return types.
+    Ignore,
+    /// Apply typed decorators and preserve the class for untyped decorators.
+    #[default]
+    IgnoreUntyped,
+    /// Apply all decorators, including those that return `Any` or `Unknown`.
+    Apply,
+}
+
 /// How to handle when recursion depth limit is exceeded.
 #[derive(Debug, PartialEq, Eq, Deserialize, Serialize, Clone, Copy, Default)]
 #[derive(ValueEnum)]
@@ -320,6 +343,9 @@ pub struct ConfigBase {
     ///   are only eligible when `check-unannotated-defs` is true.
     pub infer_return_types: Option<InferReturnTypes>,
 
+    /// How class decorators affect the class type. Defaults to `ignore-untyped`.
+    pub class_decorator_behavior: Option<ClassDecoratorBehavior>,
+
     /// Whether to disable type errors in language server. By default errors will be shown in IDEs.
     pub disable_type_errors_in_ide: Option<bool>,
 
@@ -451,6 +477,10 @@ impl ConfigBase {
         base.infer_return_types
     }
 
+    pub fn get_class_decorator_behavior(base: &Self) -> Option<ClassDecoratorBehavior> {
+        base.class_decorator_behavior
+    }
+
     pub fn get_disable_type_errors_in_ide(base: &Self) -> Option<bool> {
         base.disable_type_errors_in_ide
     }

crates/pyrefly_config/src/config.rs

@@ -67,6 +67,7 @@ use uv_pep440::Version;
 #[cfg(not(target_arch = "wasm32"))]
 use uv_pep440::VersionSpecifiers;
 
+use crate::base::ClassDecoratorBehavior;
 use crate::base::ConfigBase;
 use crate::base::ExtraConfigs;
 use crate::base::InferReturnTypes;
@@ -1213,6 +1214,15 @@ impl ConfigFile {
             .unwrap_or_else(|| self.root.infer_return_types.unwrap())
     }
 
+    pub fn class_decorator_behavior(&self, path: &ModulePath) -> ClassDecoratorBehavior {
+        self.get_from_config_overrides(ConfigBase::get_class_decorator_behavior, path)
+            .unwrap_or_else(|| {
+                self.root
+                    .class_decorator_behavior
+                    .expect("configure() sets the root class decorator behavior")
+            })
+    }
+
     pub fn disable_type_errors_in_ide(&self, path: &ModulePath) -> bool {
         self.get_from_config_overrides(ConfigBase::get_disable_type_errors_in_ide, path)
             .unwrap_or_else(|| self.root.disable_type_errors_in_ide.unwrap_or_default())
@@ -1755,6 +1765,7 @@ impl ConfigFile {
             }
             apply_preset_default!(check_unannotated_defs);
             apply_preset_default!(infer_return_types);
+            apply_preset_default!(class_decorator_behavior);
             apply_preset_default!(infer_with_first_use);
             apply_preset_default!(strict_callable_subtyping);
             apply_preset_default!(strict_partial_subtyping);
@@ -1807,6 +1818,10 @@ impl ConfigFile {
             self.root.infer_return_types = Some(InferReturnTypes::Checked);
         }
 
+        if self.root.class_decorator_behavior.is_none() {
+            self.root.class_decorator_behavior = Some(ClassDecoratorBehavior::default());
+        }
+
         if self.root.ignore_errors_in_generated_code.is_none() {
             self.root.ignore_errors_in_generated_code = Some(Default::default());
         }
@@ -2642,6 +2657,7 @@ mod tests {
                     untyped_def_behavior: Some(UntypedDefBehavior::CheckAndInferReturnType),
                     check_unannotated_defs: None,
                     infer_return_types: None,
+                    class_decorator_behavior: None,
                     permissive_ignores: None,
                     enabled_ignores: None,
                     type_ignore_unknown_tag_behavior: None,
@@ -2673,6 +2689,7 @@ mod tests {
                         untyped_def_behavior: Some(UntypedDefBehavior::CheckAndInferReturnAny),
                         check_unannotated_defs: None,
                         infer_return_types: None,
+                        class_decorator_behavior: None,
                         permissive_ignores: None,
                         enabled_ignores: None,
                         type_ignore_unknown_tag_behavior: None,
@@ -3360,6 +3377,7 @@ output-format = "omit-errors"
                 untyped_def_behavior: Some(UntypedDefBehavior::CheckAndInferReturnType),
                 check_unannotated_defs: None,
                 infer_return_types: None,
+                class_decorator_behavior: None,
                 disable_type_errors_in_ide: Some(true),
                 ignore_errors_in_generated_code: Some(false),
                 infer_with_first_use: Some(true),
@@ -3635,6 +3653,49 @@ output-format = "omit-errors"
         );
     }
 
+    #[test]
+    fn test_class_decorator_behavior_config() {
+        let mut config = ConfigFile::default();
+        config.interpreters.skip_interpreter_query = true;
+        assert!(config.configure().is_empty());
+        assert_eq!(
+            config.class_decorator_behavior(&test_path("src/main.py")),
+            ClassDecoratorBehavior::IgnoreUntyped,
+        );
+
+        for (value, expected) in [
+            ("ignore", ClassDecoratorBehavior::Ignore),
+            ("ignore-untyped", ClassDecoratorBehavior::IgnoreUntyped),
+            ("apply", ClassDecoratorBehavior::Apply),
+        ] {
+            let mut config: ConfigFile = toml::from_str(&format!(
+                "class-decorator-behavior = '{value}'\n\
+                 [[sub-config]]\n\
+                 matches = 'legacy/**'\n\
+                 class-decorator-behavior = 'ignore'\n"
+            ))
+            .unwrap();
+            config.interpreters.skip_interpreter_query = true;
+            assert!(config.configure().is_empty());
+            assert_eq!(
+                config.class_decorator_behavior(&test_path("src/main.py")),
+                expected,
+            );
+            assert_eq!(
+                config.class_decorator_behavior(&test_path("legacy/main.py")),
+                ClassDecoratorBehavior::Ignore,
+            );
+            let target_config = default_config_test_config(Some(serde_json::json!({
+                "class-decorator-behavior": value
+            })));
+            assert_eq!(
+                target_config.class_decorator_behavior(&test_path("src/main.py")),
+                expected,
+            );
+        }
+        assert!(toml::from_str::<ConfigFile>("class-decorator-behavior = 'invalid'").is_err());
+    }
+
     /// `type-ignore-unknown-tag-behavior` resolves through the same chain as every
     /// other per-file setting, so a build system can supply it. `default_config`
     /// reaches it by the same path once the target config declines to answer.
@@ -4740,6 +4801,7 @@ output-format = "omit-errors"
                 untyped_def_behavior: Some(UntypedDefBehavior::CheckAndInferReturnType),
                 check_unannotated_defs: None,
                 infer_return_types: None,
+                class_decorator_behavior: None,
                 disable_type_errors_in_ide: Some(true),
                 ignore_errors_in_generated_code: Some(false),
                 infer_with_first_use: Some(true),
@@ -4784,6 +4846,7 @@ output-format = "omit-errors"
                 untyped_def_behavior: Some(UntypedDefBehavior::CheckAndInferReturnType),
                 check_unannotated_defs: None,
                 infer_return_types: None,
+                class_decorator_behavior: None,
                 disable_type_errors_in_ide: Some(true),
                 ignore_errors_in_generated_code: Some(false),
                 infer_with_first_use: Some(true),

crates/pyrefly_types/src/stdlib.rs

@@ -323,6 +323,10 @@ impl Stdlib {
         Self::new_with_bootstrapping(true, PythonVersion::default(), &|_, _| None, &|_, _| None)
     }
 
+    pub fn is_bootstrapping(&self) -> bool {
+        self.object.is_err()
+    }
+
     fn unwrap<T>(x: &StdlibResult<T>) -> &T {
         match x {
             Ok(x) => x,

pyrefly/lib/alt/function.rs

@@ -2064,6 +2064,67 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         self.singledispatch_dispatcher_as_callback(widened, &application.original_decoratee)
     }
 
+    pub(crate) fn apply_class_decorator(
+        &self,
+        decorator: Type,
+        decoratee: Type,
+        range: TextRange,
+        errors: &ErrorCollector,
+    ) -> Type {
+        if self.stdlib.is_bootstrapping() {
+            return decoratee;
+        }
+        let preserves_class = match &decorator {
+            Type::KwCall(call) => {
+                call.func_metadata
+                    .flags
+                    .dataclass_transform_metadata
+                    .is_some()
+                    || call.has_function_kind(FunctionKind::Dataclass)
+                    || call.has_function_kind(FunctionKind::DataclassTransform)
+                    || call.has_function_kind(FunctionKind::DisjointBase)
+                    || call.has_function_kind(FunctionKind::Final)
+                    || call.has_function_kind(FunctionKind::RuntimeCheckable)
+                    || call.has_function_kind(FunctionKind::TotalOrdering)
+            }
+            _ => decorator.toplevel_func_metadata().is_some_and(|meta| {
+                meta.flags.dataclass_transform_metadata.is_some()
+                    || matches!(
+                        meta.kind,
+                        FunctionKind::Dataclass
+                            | FunctionKind::DataclassTransform
+                            | FunctionKind::DisjointBase
+                            | FunctionKind::Final
+                            | FunctionKind::RuntimeCheckable
+                            | FunctionKind::TotalOrdering
+                    )
+            }),
+        };
+        if preserves_class {
+            return decoratee;
+        }
+        let call_target =
+            self.as_call_target_or_error(decorator, CallStyle::FreeForm, range, errors, None);
+        // Keep the class's type parameters concrete during inference, just as for generic
+        // function decorators. Promoting a class definition here would replace them with defaults.
+        let (decoratee_tparams, decoratee_arg) = match &decoratee {
+            Type::ClassDef(cls) => (
+                self.get_class_tparams(cls).cloned(),
+                self.heap.mk_type_of(self.instantiate(cls)),
+            ),
+            _ => (None, decoratee.clone()),
+        };
+        let arg = CallArg::ty(&decoratee_arg, range);
+        let decorated = self
+            .call_infer(call_target, &[arg], &[], range, errors, None, None, None)
+            .ty;
+        if decorated == decoratee_arg {
+            decoratee
+        } else {
+            self.restore_decoratee_generics(decorated, decoratee_tparams)
+        }
+    }
+
     /// For a type guard function, validate whether it has at least one
     /// positional argument.
     fn validate_type_guard_positional_argument_count(