fix decorated class not transformed #1180

crates/pyrefly_config/src/args.rs

@@ -23,6 +23,7 @@ use pyrefly_util::arc_id::ArcId;
 use pyrefly_util::display;
 use pyrefly_util::interned_path::InternedPath;
 
+use crate::base::ClassDecoratorBehavior;
 use crate::base::InferReturnTypes;
 use crate::base::Preset;
 use crate::base::RecursionOverflowHandler;
@@ -370,6 +371,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, typed-only (default), or all.
+    #[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(
@@ -502,6 +506,7 @@ impl ConfigOverrideArgs {
             untyped_def_behavior,
             check_unannotated_defs,
             infer_return_types,
+            class_decorator_behavior,
             permissive_ignores,
             enabled_ignores,
             type_ignore_unknown_tag_behavior,
@@ -534,6 +539,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()
@@ -643,6 +649,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.
@@ -881,6 +890,33 @@ mod tests {
         );
     }
 
+    #[test]
+    fn class_decorator_behavior_cli_overrides_config() {
+        for (value, expected) in [
+            ("ignore", ClassDecoratorBehavior::Ignore),
+            ("typed-only", ClassDecoratorBehavior::TypedOnly),
+            ("all", ClassDecoratorBehavior::All),
+        ] {
+            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::All);
+            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

@@ -51,6 +51,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 {
+    /// Ignore inferred decorator results, but honor explicit `type[Any]` return annotations.
+    Ignore,
+    /// Apply typed decorators and preserve the class for untyped decorators.
+    #[default]
+    TypedOnly,
+    /// Apply all decorators, including those that return `Any` or `Unknown`.
+    All,
+}
+
 /// How to handle when recursion depth limit is exceeded.
 #[derive(Debug, PartialEq, Eq, Deserialize, Serialize, Clone, Copy, Default)]
 #[derive(ValueEnum)]
@@ -175,6 +198,7 @@ impl Preset {
                     check_unannotated_defs: Some(false),
                     infer_return_types: Some(InferReturnTypes::Never),
                     infer_with_first_use: Some(false),
+                    class_decorator_behavior: Some(ClassDecoratorBehavior::Ignore),
                     enabled_ignores: Some(Tool::all()),
                     ..Default::default()
                 }
@@ -194,6 +218,7 @@ impl Preset {
                     check_unannotated_defs: Some(false),
                     infer_return_types: Some(InferReturnTypes::Never),
                     legacy_overload_expansion: Some(true),
+                    class_decorator_behavior: Some(ClassDecoratorBehavior::Ignore),
                     type_ignore_unknown_tag_behavior: Some(TypeIgnoreUnknownTagBehavior::Suppress),
                     ..Default::default()
                 }
@@ -336,6 +361,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 `typed-only`.
+    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>,
 
@@ -432,6 +460,7 @@ impl ConfigBase {
             replace_untyped_imports_with_any: Some(Vec::new()),
             check_unannotated_defs: Some(true),
             infer_return_types: Some(InferReturnTypes::Checked),
+            class_decorator_behavior: Some(ClassDecoratorBehavior::default()),
             disable_type_errors_in_ide: Some(false),
             ignore_errors_in_generated_code: Some(false),
             infer_with_first_use: Some(true),
@@ -462,6 +491,7 @@ impl ConfigBase {
             untyped_def_behavior,
             check_unannotated_defs,
             infer_return_types,
+            class_decorator_behavior,
             disable_type_errors_in_ide,
             ignore_errors_in_generated_code,
             infer_with_first_use,
@@ -496,6 +526,7 @@ impl ConfigBase {
             untyped_def_behavior,
             check_unannotated_defs,
             infer_return_types,
+            class_decorator_behavior,
             disable_type_errors_in_ide,
             ignore_errors_in_generated_code,
             infer_with_first_use,

crates/pyrefly_config/src/config.rs

@@ -69,6 +69,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;
@@ -1352,6 +1353,13 @@ impl ConfigFile {
             .expect(RESOLVED_SETTING)
     }
 
+    pub fn class_decorator_behavior(&self, path: &ModulePath) -> ClassDecoratorBehavior {
+        self.resolved(path)
+            .settings
+            .class_decorator_behavior
+            .expect(RESOLVED_SETTING)
+    }
+
     pub fn infer_with_first_use(&self, path: &ModulePath) -> bool {
         self.resolved(path)
             .settings
@@ -2906,6 +2914,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,
@@ -2942,6 +2951,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,
@@ -3678,6 +3688,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),
@@ -3988,6 +3999,77 @@ 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::TypedOnly,
+        );
+
+        for (value, expected) in [
+            ("ignore", ClassDecoratorBehavior::Ignore),
+            ("typed-only", ClassDecoratorBehavior::TypedOnly),
+            ("all", ClassDecoratorBehavior::All),
+        ] {
+            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());
+    }
+
+    #[test]
+    fn test_class_decorator_behavior_presets() {
+        for (preset, expected) in [
+            (Preset::Basic, ClassDecoratorBehavior::Ignore),
+            (Preset::Legacy, ClassDecoratorBehavior::Ignore),
+            (Preset::Default, ClassDecoratorBehavior::TypedOnly),
+            (Preset::Strict, ClassDecoratorBehavior::TypedOnly),
+            (Preset::All, ClassDecoratorBehavior::TypedOnly),
+        ] {
+            for override_behavior in [
+                None,
+                Some(ClassDecoratorBehavior::Ignore),
+                Some(ClassDecoratorBehavior::TypedOnly),
+                Some(ClassDecoratorBehavior::All),
+            ] {
+                let mut config = ConfigFile::default();
+                config.interpreters.skip_interpreter_query = true;
+                config.root.preset = Some(preset);
+                config.root.class_decorator_behavior = override_behavior;
+                assert!(config.configure().is_empty());
+                assert_eq!(
+                    config.class_decorator_behavior(&test_path("src/main.py")),
+                    override_behavior.unwrap_or(expected),
+                );
+            }
+        }
+    }
+
     /// `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.
@@ -5405,6 +5487,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),
@@ -5451,6 +5534,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

@@ -318,6 +318,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/call.rs

@@ -3094,6 +3094,44 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         }
     }
 
+    /// Preserve unknown `__new__` returns when a constructor is used as a class decorator.
+    pub(crate) fn class_decorator_constructor_to_callable(&self, cls: &ClassType) -> Type {
+        let mut constructor = self.constructor_to_callable(cls);
+        if let Some(new_member) =
+            self.get_class_member_with_defining_class(cls.class_object(), &dunder::NEW)
+            && new_member.value.is_function_without_return_annotation()
+            && let Some(new_metadata) = new_member.value.ty().toplevel_func_metadata()
+        {
+            // Ordinary construction gives an unannotated `__new__` a synthetic `Self` return.
+            // That fallback does not describe a decorator that returns its class argument.
+            constructor = self.unions(
+                constructor
+                    .into_unions()
+                    .into_iter()
+                    .map(|mut member| {
+                        if member
+                            .toplevel_func_metadata()
+                            .is_some_and(|meta| meta.kind == new_metadata.kind)
+                        {
+                            member.transform_toplevel_callable_signatures(|signature, _| {
+                                signature.ret = self.heap.mk_any_implicit();
+                            });
+                        }
+                        member
+                    })
+                    .collect(),
+            );
+        }
+        if let Some(tparams) = self
+            .get_class_tparams(cls.class_object())
+            .filter(|p| !p.is_empty())
+        {
+            self.normalize_class_constructor_tparams(constructor, tparams)
+        } else {
+            constructor
+        }
+    }
+
     /// Convert a bare class definition while keeping its type parameters generic.
     pub fn constructor_to_callable_for_class_def(&self, cls: &Class) -> Type {
         let Some(class_tparams) = self

pyrefly/lib/alt/class/class_field.rs

@@ -1242,7 +1242,7 @@ impl ClassField {
         }
     }
 
-    fn is_function_without_return_annotation(&self) -> bool {
+    pub(crate) fn is_function_without_return_annotation(&self) -> bool {
         match &self.0 {
             ClassFieldInner::Property { .. } => false,
             ClassFieldInner::Descriptor { .. } => false,