From a44977f341d5900c8f35302c287bc4dd8908cfe8 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 24 Oct 2026 02:01:35 +0900 Subject: [PATCH 1/6] test --- pyrefly/lib/test/decorators.rs | 129 +++++++++++++++++++++++++++++++++ pyrefly/lib/test/shape_dsl.rs | 58 +++++++++++++++ 2 files changed, 187 insertions(+) diff --git a/pyrefly/lib/test/decorators.rs b/pyrefly/lib/test/decorators.rs index 7e5033a813..59f56813e5 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -1120,6 +1120,135 @@ g(f) "#, ); +testcase!( + test_class_decorator_return_type_applied, + r#" +from typing import reveal_type + +def to_int(cls: type) -> int: + return 42 + +@to_int +class C: + ... + +reveal_type(C) # E: revealed type: int + "#, +); + +testcase!( + test_class_decorator_dataclass_preserves_class, + r#" +from dataclasses import dataclass +from typing import assert_type + +@dataclass +class Bare: + value: int + +@dataclass(kw_only=True) +class Called: + value: int + +assert_type(Bare, type[Bare]) +assert_type(Bare(1), Bare) +assert_type(Bare(1).value, int) +assert_type(Called, type[Called]) +assert_type(Called(value=1), Called) +assert_type(Called(value=1).value, int) +Called(1) # E: Expected argument `value` to be passed by name + "#, +); + +testcase!( + test_class_decorator_dataclass_transform_preserves_class, + r#" +from typing import assert_type, dataclass_transform + +@dataclass_transform(kw_only_default=True) +class Base: ... + +class Child(Base): + value: int + +assert_type(Base, type[Base]) +assert_type(Base(), Base) +assert_type(Child(value=1), Child) +assert_type(Child(value=1).value, int) +Child(1) # E: Expected argument `value` to be passed by name + "#, +); + +testcase!( + test_class_decorator_dataclass_transform_metadata_preserves_class, + r#" +from typing import Callable, assert_type, dataclass_transform + +# Transform metadata preserves the class even when the signature erases its type. +@dataclass_transform() +def model(cls: type[object]) -> type[object]: ... + +@dataclass_transform() +def model_factory(*, kw_only: bool = False) -> Callable[[type[object]], type[object]]: ... + +@model +class Bare: + value: int + +@model_factory(kw_only=True) +class Called: + value: int + +assert_type(Bare, type[Bare]) +assert_type(Bare(1), Bare) +assert_type(Bare(1).value, int) +assert_type(Called, type[Called]) +assert_type(Called(value=1), Called) +assert_type(Called(value=1).value, int) +Called(1) # E: Expected argument `value` to be passed by name + "#, +); + +testcase!( + test_class_decorator_special_function_kinds_preserve_class, + r#" +from functools import total_ordering +from typing import Protocol, assert_type, final, runtime_checkable +from typing_extensions import disjoint_base + +@final +class FinalClass: + value: int = 1 + +@disjoint_base +class DisjointClass: + value: int = 1 + +@runtime_checkable +class RuntimeProtocol(Protocol): + value: int + +@total_ordering +class OrderedClass: + def __lt__(self, other: OrderedClass) -> bool: ... + +assert_type(FinalClass, type[FinalClass]) +assert_type(FinalClass(), FinalClass) +assert_type(FinalClass().value, int) +assert_type(DisjointClass, type[DisjointClass]) +assert_type(DisjointClass(), DisjointClass) +assert_type(DisjointClass().value, int) +assert_type(OrderedClass, type[OrderedClass]) +assert_type(OrderedClass(), OrderedClass) +assert_type(OrderedClass() >= OrderedClass(), bool) + +def narrow(value: object) -> None: + if isinstance(value, RuntimeProtocol): + assert_type(value, RuntimeProtocol) + assert_type(value.value, int) + "#, +); + testcase!( test_decorator_missing_concatenate_parameters, r#" diff --git a/pyrefly/lib/test/shape_dsl.rs b/pyrefly/lib/test/shape_dsl.rs index 2aded51850..dfcbd8e311 100644 --- a/pyrefly/lib/test/shape_dsl.rs +++ b/pyrefly/lib/test/shape_dsl.rs @@ -9398,6 +9398,64 @@ def named_variadic(x: Float[Array, "*batch channels"]) -> None: mod legacy { use super::*; + testcase!( + test_shaped_array_decorator_preserves_class, + legacy_shaped_array_env(), + r#" +import shape_extensions as se +from shape_extensions import IntTuple, shaped_array +from shape_extensions import shaped_array as array_metadata +from typing import assert_type + +@shaped_array(shape="Shape") +class Direct[Shape: IntTuple]: + value: int = 1 + +@array_metadata(shape="Shape") +class Alias[Shape: IntTuple]: + value: int = 1 + +@se.shaped_array(shape="Shape") +class Qualified[Shape: IntTuple]: + value: int = 1 + +assert_type(Direct().value, int) +assert_type(Alias().value, int) +assert_type(Qualified().value, int) + +def check_shapes( + direct: Direct[[2, 3]], alias: Alias[[2, 3]], qualified: Qualified[[2, 3]] +) -> None: + assert_type(direct[0], Direct[[3]]) + assert_type(alias[0], Alias[[3]]) + assert_type(qualified[0], Qualified[[3]]) + +# Invalid bare metadata decorators still preserve the class during recovery. +@shaped_array # E: `@shaped_array` requires a `shape` keyword argument +class Bare: + value: int = 1 + +assert_type(Bare, type[Bare]) +assert_type(Bare().value, int) +"#, + ); + + testcase!( + test_shaped_array_decorator_name_is_not_metadata, + legacy_shaped_array_env(), + r#" +from typing import Callable, assert_type + +# Only the special export is metadata, not an unrelated function with the same name. +def shaped_array(*, shape: str) -> Callable[[type], int]: ... + +@shaped_array(shape="Shape") +class Ordinary: ... + +assert_type(Ordinary, int) +"#, + ); + #[test] fn test_shaped_array_imports_are_metadata() { let mut env = legacy_shaped_array_env(); From 41caf5f314723a67f308cc32e1eca69a9bdd9086 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Thu, 8 Oct 2026 11:30:47 +0900 Subject: [PATCH 2/6] Apply class decorator return types while preserving special decorators --- crates/pyrefly_types/src/stdlib.rs | 4 +++ pyrefly/lib/alt/function.rs | 46 ++++++++++++++++++++++++ pyrefly/lib/alt/solve.rs | 58 +++++++++++++++--------------- pyrefly/lib/binding/binding.rs | 4 ++- pyrefly/lib/binding/expr.rs | 7 ++++ 5 files changed, 90 insertions(+), 29 deletions(-) diff --git a/crates/pyrefly_types/src/stdlib.rs b/crates/pyrefly_types/src/stdlib.rs index f4ded04b50..988eae1339 100644 --- a/crates/pyrefly_types/src/stdlib.rs +++ b/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(x: &StdlibResult) -> &T { match x { Ok(x) => x, diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index 2927c3c510..afd00585f2 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -2300,6 +2300,52 @@ 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); + let arg = CallArg::ty(&decoratee, range); + self.call_infer(call_target, &[arg], &[], range, errors, None, None, None) + .ty + } + /// For a type guard function, validate whether it has at least one /// positional argument. fn validate_type_guard_positional_argument_count( diff --git a/pyrefly/lib/alt/solve.rs b/pyrefly/lib/alt/solve.rs index 454f15cd24..85a74ba269 100644 --- a/pyrefly/lib/alt/solve.rs +++ b/pyrefly/lib/alt/solve.rs @@ -7518,35 +7518,37 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ); } } - // TODO: analyze the class decorators beyond the unknown-type check above. We don't - // support general type-level analysis of class decorators (the ones we do - // support, like dataclass-related ones, are handled via custom bindings). - // - // Note that all decorators have their own binding so they are still type checked for errors - // *inside* the decorator. The only application-level effect we honor is an - // explicit `type[Any]` return, which intentionally erases the class interface. - // - // We use `.any()` across the chain because if *any* decorator is annotated - // with `-> type[Any]`, all *later* decorators receive whatever class that - // dynamic decorator produces (an `Any`-ish class), so the final output must - // also be treated as dynamic. This is by intent: the escape hatch lets library - // authors signal that the class flows through a decorator a type checker cannot - // model accurately, so the end result is not modelable either. - let erases_class = decorators.iter().any(|idx| { - let decorator_ty = &self.get_idx(*idx).ty; - // Only an explicit `-> type[Any]` return annotation counts; a return - // type inferred from the body (e.g. `return cls` where `cls: type[Any]`) - // does not, since the user did not opt in to erasing the class. - !decorator_ty.toplevel_func_metadata().is_some_and(|meta| meta.flags.is_return_inferred) - && decorator_ty - .toplevel_callable_signatures() - .any(|(c, _)| matches!(&c.ret, Type::Type(inner) if matches!(&**inner, Type::Any(AnyStyle::Explicit)))) - }); - if erases_class { - self.heap.mk_type_of(self.heap.mk_any_explicit()) - } else { - self.heap.mk_class_def(cls.dupe()) + let mut ty = self.heap.mk_class_def(cls.dupe()); + for decorator_key in decorators.iter().rev() { + if self.bindings().get(*decorator_key).is_class_metadata { + continue; + } + let decorator = self.get_idx(*decorator_key); + let preserves_class = matches!( + &decorator.ty, + Type::Any(AnyStyle::Implicit | AnyStyle::Explicit) + ) || decorator.ty.toplevel_callable_signatures().any( + |(callable, _)| { + matches!(&callable.ret, Type::Any(_)) + || decorator.ty.toplevel_func_metadata().is_some_and(|meta| { + meta.flags.is_return_inferred + && matches!( + &callable.ret, + Type::Type(inner) if matches!( + &**inner, + Type::Any(AnyStyle::Explicit) + ) + ) + }) + }, + ); + if preserves_class { + continue; + } + let range = self.bindings().idx_to_key(*decorator_key).range(); + ty = self.apply_class_decorator(decorator.ty.clone(), ty, range, errors); } + ty } }, Binding::AnnotatedType(ann, val) => { diff --git a/pyrefly/lib/binding/binding.rs b/pyrefly/lib/binding/binding.rs index 890e14c7df..21c99330ef 100644 --- a/pyrefly/lib/binding/binding.rs +++ b/pyrefly/lib/binding/binding.rs @@ -149,7 +149,7 @@ assert_words!(BindingClassSynthesizedFields, 2); assert_bytes!(BindingLegacyTypeParam, 16); assert_words!(BindingYield, 4); assert_words!(BindingYieldFrom, 4); -assert_words!(BindingDecorator, 9); +assert_words!(BindingDecorator, 10); assert_bytes!(BindingDecoratedFunction, 24); assert_words!(BindingUndecoratedFunction, 18); @@ -2092,6 +2092,8 @@ pub enum FunctionParameter { pub struct BindingDecorator { pub expr: Expr, pub trailing_name: Option, + /// Metadata decorators are consumed during binding and must not replace the class object. + pub is_class_metadata: bool, } impl DisplayWith for BindingDecorator { diff --git a/pyrefly/lib/binding/expr.rs b/pyrefly/lib/binding/expr.rs index 8eda8091f8..498a29db3f 100644 --- a/pyrefly/lib/binding/expr.rs +++ b/pyrefly/lib/binding/expr.rs @@ -1712,6 +1712,12 @@ impl<'a> BindingsBuilder<'a> { ) -> Vec> { let mut decorator_keys = Vec::with_capacity(decorators.len()); for mut x in decorators { + let decorator_func = x + .expression + .as_call_expr() + .map_or(&x.expression, |call| &call.func); + let is_class_metadata = + self.as_special_export(decorator_func) == Some(SpecialExport::ShapedArray); self.ensure_expr(&mut x.expression, usage); let trailing_name = Ast::decorator_trailing_name(&x.expression).map(Name::new); let k = self.insert_binding( @@ -1719,6 +1725,7 @@ impl<'a> BindingsBuilder<'a> { BindingDecorator { expr: x.expression, trailing_name, + is_class_metadata, }, ); decorator_keys.push(k); From 8e7e95eec2ff4ce55282438c6be3870a8699b830 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 24 Oct 2026 17:24:52 +0900 Subject: [PATCH 3/6] fix test --- pyrefly/lib/alt/function.rs | 17 ++++++++++++-- pyrefly/lib/alt/solve.rs | 41 +++++++++++++++++---------------- pyrefly/lib/binding/bindings.rs | 12 ++++++++++ pyrefly/lib/test/calls.rs | 5 ++++ pyrefly/lib/test/decorators.rs | 35 ++++++++++++++++++++++++++++ 5 files changed, 88 insertions(+), 22 deletions(-) diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index afd00585f2..d2de37cd2d 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -2342,8 +2342,21 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let call_target = self.as_call_target_or_error(decorator, CallStyle::FreeForm, range, errors, None); let arg = CallArg::ty(&decoratee, range); - self.call_infer(call_target, &[arg], &[], range, errors, None, None, None) - .ty + let decorated = self + .call_infer(call_target, &[arg], &[], range, errors, None, None, None) + .ty; + // Retain class-definition metadata when the decorator returns the same nongeneric class. + // A specialized generic return must keep its inferred type arguments. + if let Type::ClassDef(cls) = &decoratee + && let Type::Type(inner) = &decorated + && let Type::ClassType(result) = &**inner + && result.class_object() == cls + && result.targs().is_empty() + { + decoratee + } else { + decorated + } } /// For a type guard function, validate whether it has at least one diff --git a/pyrefly/lib/alt/solve.rs b/pyrefly/lib/alt/solve.rs index 85a74ba269..d8ce3bb9d2 100644 --- a/pyrefly/lib/alt/solve.rs +++ b/pyrefly/lib/alt/solve.rs @@ -7481,8 +7481,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { None => self.heap.mk_any_implicit(), Some(cls) => { // A class decorator obscures the class if its own type is `Any`, or if it - // returns an unknown type. Class decorators are not applied below, so the - // result is read from the decorator's signatures rather than by calling it. + // returns an unknown type. Read the return signatures before applying the + // decorators so this diagnostic does not depend on application results. // Like for functions, only the first decorator (in application order) that // makes the result unknown is blamed for it. let returns_unknown = |callable: &Type| { @@ -7524,24 +7524,25 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { continue; } let decorator = self.get_idx(*decorator_key); - let preserves_class = matches!( - &decorator.ty, - Type::Any(AnyStyle::Implicit | AnyStyle::Explicit) - ) || decorator.ty.toplevel_callable_signatures().any( - |(callable, _)| { - matches!(&callable.ret, Type::Any(_)) - || decorator.ty.toplevel_func_metadata().is_some_and(|meta| { - meta.flags.is_return_inferred - && matches!( - &callable.ret, - Type::Type(inner) if matches!( - &**inner, - Type::Any(AnyStyle::Explicit) - ) - ) - }) - }, - ); + let preserves_class = matches!(&decorator.ty, Type::Any(_)) + || decorator + .ty + .toplevel_callable_signatures() + .any(|(callable, _)| { + matches!(&callable.ret, Type::Any(_)) + || decorator.ty.toplevel_func_metadata().is_some_and( + |meta| { + meta.flags.is_return_inferred + && matches!( + &callable.ret, + Type::Type(inner) if matches!( + &**inner, + Type::Any(AnyStyle::Explicit) + ) + ) + }, + ) + }); if preserves_class { continue; } diff --git a/pyrefly/lib/binding/bindings.rs b/pyrefly/lib/binding/bindings.rs index 86965f6a74..e5a88011e0 100644 --- a/pyrefly/lib/binding/bindings.rs +++ b/pyrefly/lib/binding/bindings.rs @@ -2687,6 +2687,18 @@ impl<'a> BindingsBuilder<'a> { let binding = if forward_to_first_use { Binding::ForwardToFirstUse(target_idx) + } else if deferred.usage.is_static() + && let Some(Binding::ClassDef(class_idx, decorators)) = self.idx_to_binding(target_idx) + && !decorators.is_empty() + && let Some(BindingClass::ClassDef(class)) = self.idx_to_binding(*class_idx) + && class + .def + .range + .contains_range(self.idx_to_key(deferred.bound_name_idx).range()) + { + // Self-references in class annotations refer to the undecorated class. Depending on + // its decorated value would create a cycle when a decorator's signature uses it. + Binding::ClassDef(*class_idx, Box::new([])) } else if let Some(b) = self.idx_to_binding(target_idx) && matches!( b, diff --git a/pyrefly/lib/test/calls.rs b/pyrefly/lib/test/calls.rs index d0d7b02749..8731deac9e 100644 --- a/pyrefly/lib/test/calls.rs +++ b/pyrefly/lib/test/calls.rs @@ -949,6 +949,11 @@ def identity[T](cls: type[T]) -> type[T]: class Old: ... Old() # E: `Old` is deprecated +alias = Old # E: `Old` is deprecated +alias() # E: `Old` is deprecated + +def accepts_subclass(cls: type[Old]) -> None: # E: `Old` is deprecated + cls() "#, ); diff --git a/pyrefly/lib/test/decorators.rs b/pyrefly/lib/test/decorators.rs index 59f56813e5..04bcc3326c 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -1136,6 +1136,41 @@ reveal_type(C) # E: revealed type: int "#, ); +testcase!( + test_class_decorator_error_preserves_inheritance, + r#" +from typing import assert_type + +@missing_decorator # E: Could not find name `missing_decorator` +class Base: + def inherited(self) -> int: + return 1 + +class Child(Base): ... + +assert_type(Base(), Base) +assert_type(Child().inherited(), int) + "#, +); + +testcase!( + test_class_decorator_preserves_returned_type_arguments, + r#" +from typing import Callable, assert_type, cast + +def identity[T](cls: type[T]) -> type[T]: + return cls + +@cast("Callable[[type[Box]], type[Box[int]]]", identity) +class Box[T]: + def __init__(self, value: T) -> None: + self.value = value + +assert_type(Box(1).value, int) +Box("bad") # E: is not assignable to parameter `value` with type `int` + "#, +); + testcase!( test_class_decorator_dataclass_preserves_class, r#" From f9e270f588041371bfcd07dd71f94902ddec58c5 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 30 Oct 2026 00:51:02 +0900 Subject: [PATCH 4/6] 3 config --- crates/pyrefly_config/src/args.rs | 36 +++++++ crates/pyrefly_config/src/base.rs | 30 ++++++ crates/pyrefly_config/src/config.rs | 63 +++++++++++ pyrefly/lib/alt/solve.rs | 57 ++++++---- pyrefly/lib/solver/solver.rs | 3 + pyrefly/lib/state/state.rs | 3 + pyrefly/lib/state/steps.rs | 3 + pyrefly/lib/test/decorators.rs | 156 ++++++++++++++++++++++++++++ pyrefly/lib/test/util.rs | 9 ++ schemas/pyrefly.json | 6 ++ schemas/validate_schemas.py | 15 +++ website/docs/configuration.mdx | 35 +++++++ 12 files changed, 395 insertions(+), 21 deletions(-) diff --git a/crates/pyrefly_config/src/args.rs b/crates/pyrefly_config/src/args.rs index 5f40e370e8..771b5e9706 100644 --- a/crates/pyrefly_config/src/args.rs +++ b/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, + /// How class decorators affect class types: ignore, typed-only (default), or all. + #[arg(long)] + class_decorator_behavior: Option, /// 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( @@ -450,6 +454,7 @@ impl ConfigOverrideArgs { untyped_def_behavior, check_unannotated_defs, infer_return_types, + class_decorator_behavior, permissive_ignores, enabled_ignores, type_ignore_unknown_tag_behavior, @@ -480,6 +485,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() @@ -585,6 +591,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. @@ -785,6 +794,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 [ diff --git a/crates/pyrefly_config/src/base.rs b/crates/pyrefly_config/src/base.rs index 1bddfc393f..b4e83a9aa6 100644 --- a/crates/pyrefly_config/src/base.rs +++ b/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] + 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)] @@ -331,6 +354,9 @@ pub struct ConfigBase { /// are only eligible when `check-unannotated-defs` is true. pub infer_return_types: Option, + /// How class decorators affect the class type. Defaults to `typed-only`. + pub class_decorator_behavior: Option, + /// Whether to disable type errors in language server. By default errors will be shown in IDEs. pub disable_type_errors_in_ide: Option, @@ -467,6 +493,10 @@ impl ConfigBase { base.infer_return_types } + pub fn get_class_decorator_behavior(base: &Self) -> Option { + base.class_decorator_behavior + } + pub fn get_disable_type_errors_in_ide(base: &Self) -> Option { base.disable_type_errors_in_ide } diff --git a/crates/pyrefly_config/src/config.rs b/crates/pyrefly_config/src/config.rs index cc59ae2149..1bd4243e22 100644 --- a/crates/pyrefly_config/src/config.rs +++ b/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; @@ -1268,6 +1269,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()) @@ -1818,6 +1828,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); @@ -1871,6 +1882,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()); } @@ -2810,6 +2825,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, @@ -2842,6 +2858,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, @@ -3532,6 +3549,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), @@ -3808,6 +3826,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::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::("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. @@ -4927,6 +4988,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), @@ -4972,6 +5034,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), diff --git a/pyrefly/lib/alt/solve.rs b/pyrefly/lib/alt/solve.rs index d8ce3bb9d2..97ae862145 100644 --- a/pyrefly/lib/alt/solve.rs +++ b/pyrefly/lib/alt/solve.rs @@ -148,6 +148,7 @@ use crate::binding::narrow::NarrowOp; use crate::binding::narrow::identifier_and_chain_for_expr; use crate::binding::narrow::identifier_and_chain_prefix_for_expr; use crate::binding::shape_type::BoundShapeComment; +use crate::config::base::ClassDecoratorBehavior; use crate::config::error_kind::ErrorKind; use crate::error::collector::ErrorCollector; use crate::error::context::ErrorContext; @@ -7519,33 +7520,47 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } } let mut ty = self.heap.mk_class_def(cls.dupe()); + let behavior = self.solver().config.class_decorator_behavior; for decorator_key in decorators.iter().rev() { - if self.bindings().get(*decorator_key).is_class_metadata { + if behavior == ClassDecoratorBehavior::Ignore + || self.bindings().get(*decorator_key).is_class_metadata + { continue; } let decorator = self.get_idx(*decorator_key); - let preserves_class = matches!(&decorator.ty, Type::Any(_)) - || decorator - .ty - .toplevel_callable_signatures() - .any(|(callable, _)| { - matches!(&callable.ret, Type::Any(_)) - || decorator.ty.toplevel_func_metadata().is_some_and( - |meta| { - meta.flags.is_return_inferred - && matches!( - &callable.ret, - Type::Type(inner) if matches!( - &**inner, - Type::Any(AnyStyle::Explicit) - ) - ) - }, - ) - }); - if preserves_class { + // An invalid decorator must not erase the class during error recovery. + if matches!(&decorator.ty, Type::Any(AnyStyle::Error)) { continue; } + if behavior == ClassDecoratorBehavior::TypedOnly { + let instance_call = match &decorator.ty { + Type::ClassType(instance) => self.instance_as_dunder_call(instance), + _ => None, + }; + let callable = instance_call.as_ref().unwrap_or(&decorator.ty); + let preserves_class = matches!(callable, Type::Any(_)) + || callable + .toplevel_callable_signatures() + .any(|(signature, _)| { + matches!(&signature.ret, Type::Any(_)) + || signature.ret.has_top_level_implicit_any() + || callable.toplevel_func_metadata().is_some_and( + |meta| { + meta.flags.is_return_inferred + && matches!( + &signature.ret, + Type::Type(inner) if matches!( + &**inner, + Type::Any(AnyStyle::Explicit) + ) + ) + }, + ) + }); + if preserves_class { + continue; + } + } let range = self.bindings().idx_to_key(*decorator_key).range(); ty = self.apply_class_decorator(decorator.ty.clone(), ty, range, errors); } diff --git a/pyrefly/lib/solver/solver.rs b/pyrefly/lib/solver/solver.rs index e86ed7be8d..e2126b17fb 100644 --- a/pyrefly/lib/solver/solver.rs +++ b/pyrefly/lib/solver/solver.rs @@ -52,6 +52,7 @@ use vec1::Vec1; use crate::alt::answers::LookupAnswer; use crate::alt::answers_solver::AnswersSolver; use crate::alt::attr::AttrSubsetError; +use crate::config::base::ClassDecoratorBehavior; use crate::config::error_kind::ErrorKind; use crate::error::collector::ErrorBuilder; use crate::error::collector::ErrorCollector; @@ -559,6 +560,7 @@ struct VarState { #[derive(Debug, Clone, Copy, Default)] pub struct SolverConfig { + pub class_decorator_behavior: ClassDecoratorBehavior, pub infer_with_first_use: bool, pub tensor_shapes: bool, pub strict_callable_subtyping: bool, @@ -4551,6 +4553,7 @@ mod tests { #[test] fn finishing_takes_the_restriction_record_off_the_answer() { let solver = Solver::new(SolverConfig { + class_decorator_behavior: ClassDecoratorBehavior::default(), infer_with_first_use: true, tensor_shapes: false, strict_callable_subtyping: false, diff --git a/pyrefly/lib/state/state.rs b/pyrefly/lib/state/state.rs index 0d2d395771..01ea767970 100644 --- a/pyrefly/lib/state/state.rs +++ b/pyrefly/lib/state/state.rs @@ -1620,6 +1620,8 @@ impl<'a> Transaction<'a> { lookup: &self.lookup(module_data), check_unannotated_defs: config.check_unannotated_defs(module_data.handle.path()), infer_return_types: config.infer_return_types(module_data.handle.path()), + class_decorator_behavior: config + .class_decorator_behavior(module_data.handle.path()), infer_with_first_use: config.infer_with_first_use(module_data.handle.path()), tensor_shapes, strict_callable_subtyping: config @@ -2684,6 +2686,7 @@ impl<'a> Transaction<'a> { lookup: &self.lookup(m), check_unannotated_defs: config.check_unannotated_defs(m.handle.path()), infer_return_types: config.infer_return_types(m.handle.path()), + class_decorator_behavior: config.class_decorator_behavior(m.handle.path()), infer_with_first_use: config.infer_with_first_use(m.handle.path()), // This is a one-shot timing/diagnostic dump, so we intentionally do not // store the bit on `module_data` (no later dirty.find() re-check applies). diff --git a/pyrefly/lib/state/steps.rs b/pyrefly/lib/state/steps.rs index 5055ee3ea0..cccf2f0923 100644 --- a/pyrefly/lib/state/steps.rs +++ b/pyrefly/lib/state/steps.rs @@ -26,6 +26,7 @@ use crate::alt::answers::Answers; use crate::alt::answers::LookupAnswer; use crate::alt::answers::Solutions; use crate::binding::bindings::Bindings; +use crate::config::base::ClassDecoratorBehavior; use crate::config::base::InferReturnTypes; use crate::config::base::RecursionLimitConfig; use crate::error::collector::ErrorCollector; @@ -63,6 +64,7 @@ pub struct Context<'a, Lookup> { pub lookup: &'a Lookup, pub check_unannotated_defs: bool, pub infer_return_types: InferReturnTypes, + pub class_decorator_behavior: ClassDecoratorBehavior, pub infer_with_first_use: bool, pub tensor_shapes: bool, pub strict_callable_subtyping: bool, @@ -554,6 +556,7 @@ impl Step { exports: Arc, ) -> Arc { let solver = Solver::new(SolverConfig { + class_decorator_behavior: ctx.class_decorator_behavior, infer_with_first_use: ctx.infer_with_first_use, tensor_shapes: ctx.tensor_shapes, strict_callable_subtyping: ctx.strict_callable_subtyping, diff --git a/pyrefly/lib/test/decorators.rs b/pyrefly/lib/test/decorators.rs index 04bcc3326c..d8f52103d5 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -7,6 +7,7 @@ use pyrefly_python::sys_info::PythonVersion; +use crate::config::base::ClassDecoratorBehavior; use crate::test::util::TestEnv; use crate::testcase; @@ -1120,6 +1121,161 @@ g(f) "#, ); +testcase!( + test_class_decorator_behavior_ignore, + TestEnv::new().with_class_decorator_behavior(ClassDecoratorBehavior::Ignore), + r#" +from dataclasses import dataclass +from typing import Any, assert_type + +def typed(cls: type) -> int: ... +def untyped(cls): return cls +def erased(cls: type) -> type[Any]: ... + +@typed +class Typed: ... + +@untyped +class Untyped: ... + +@erased +class Erased: ... + +assert_type(Typed, type[Typed]) +assert_type(Untyped, type[Untyped]) +assert_type(Erased, type[Erased]) +assert_type(Typed(), Typed) + +# Special decorators still provide their synthesized class behavior. +@dataclass +class Data: + value: int + +assert_type(Data(1).value, int) +"#, +); + +testcase!( + test_class_decorator_behavior_typed_only, + TestEnv::new().with_class_decorator_behavior(ClassDecoratorBehavior::TypedOnly), + r#" +from typing import Any, Callable, assert_type + +def typed(cls: type) -> int: ... +def untyped(cls): return cls +def explicit_any(cls: type) -> Any: ... +def inferred_type_any(cls: type[Any]): return cls +def factory() -> Callable: return untyped + +class UntypedDecorator: + def __call__(self, cls): return cls + +@typed +class Typed: ... + +@untyped +class Untyped: ... + +@explicit_any +class ExplicitAny: ... + +@inferred_type_any +class InferredTypeAny: ... + +@factory() +class Factory: ... + +@UntypedDecorator() +class Instance: ... + +assert_type(Typed, int) +assert_type(Untyped, type[Untyped]) +assert_type(ExplicitAny, type[ExplicitAny]) +assert_type(InferredTypeAny, type[InferredTypeAny]) +assert_type(Factory, type[Factory]) +assert_type(Instance, type[Instance]) +"#, +); + +testcase!( + test_class_decorator_behavior_default, + r#" +from typing import assert_type + +def typed(cls: type) -> int: ... +def untyped(cls): return cls + +@typed +class Typed: ... + +@untyped +class Untyped: ... + +assert_type(Typed, int) +assert_type(Untyped, type[Untyped]) +"#, +); + +testcase!( + test_class_decorator_behavior_all, + TestEnv::new().with_class_decorator_behavior(ClassDecoratorBehavior::All), + r#" +from dataclasses import dataclass +from typing import Any, Callable, assert_type + +def typed(cls: type) -> int: ... +def untyped(cls): return cls +def explicit_any(cls: type) -> Any: ... +def inferred_type_any(cls: type[Any]): return cls +def factory() -> Callable: return untyped +def to_str(value: int) -> str: ... + +class UntypedDecorator: + def __call__(self, cls): return cls + +@typed +class Typed: ... + +@untyped +class Untyped: ... + +@explicit_any +class ExplicitAny: ... + +@inferred_type_any +class InferredTypeAny: ... + +@factory() +class Factory: ... + +@UntypedDecorator() +class Instance: ... + +@to_str +@typed +class Stacked: ... + +assert_type(Typed, int) +assert_type(Untyped, Any) +assert_type(ExplicitAny, Any) +assert_type(InferredTypeAny, type[Any]) +assert_type(Factory, Any) +assert_type(Instance, Any) +assert_type(Stacked, str) + +@dataclass(kw_only=True) +class Data: + value: int + +assert_type(Data(value=1).value, int) + +@missing_decorator # E: Could not find name `missing_decorator` +class Invalid: ... + +assert_type(Invalid(), Invalid) +"#, +); + testcase!( test_class_decorator_return_type_applied, r#" diff --git a/pyrefly/lib/test/util.rs b/pyrefly/lib/test/util.rs index 1567a125e6..c8fa382583 100644 --- a/pyrefly/lib/test/util.rs +++ b/pyrefly/lib/test/util.rs @@ -47,6 +47,7 @@ use ruff_text_size::TextSize; use tempfile::TempDir; use crate::binding::binding::KeyExport; +use crate::config::base::ClassDecoratorBehavior; use crate::config::base::InferReturnTypes; use crate::config::base::UntypedDefBehavior; use crate::config::config::ConfigFile; @@ -252,6 +253,7 @@ pub struct TestEnv { platform: PythonPlatform, check_unannotated_defs: bool, infer_return_types: InferReturnTypes, + class_decorator_behavior: ClassDecoratorBehavior, infer_with_first_use: bool, non_exhaustive_match_open_type_error: bool, recursion_depth_limit: Option, @@ -332,6 +334,7 @@ impl TestEnv { platform: PythonPlatform::default(), check_unannotated_defs: true, infer_return_types: InferReturnTypes::Checked, + class_decorator_behavior: ClassDecoratorBehavior::default(), infer_with_first_use: true, non_exhaustive_match_open_type_error: false, recursion_depth_limit: None, @@ -486,6 +489,11 @@ impl TestEnv { res } + pub fn with_class_decorator_behavior(mut self, behavior: ClassDecoratorBehavior) -> Self { + self.class_decorator_behavior = behavior; + self + } + pub fn with_recursion_depth_limit(mut self, recursion_depth_limit: u32) -> Self { self.recursion_depth_limit = Some(recursion_depth_limit); self @@ -878,6 +886,7 @@ impl TestEnv { config.python_environment.site_package_path = Some(self.site_package_path.clone()); config.root.check_unannotated_defs = Some(self.check_unannotated_defs); config.root.infer_return_types = Some(self.infer_return_types); + config.root.class_decorator_behavior = Some(self.class_decorator_behavior); config.root.infer_with_first_use = Some(self.infer_with_first_use); config.root.recursion_depth_limit = self.recursion_depth_limit; config.root.strict_callable_subtyping = Some(self.strict_callable_subtyping); diff --git a/schemas/pyrefly.json b/schemas/pyrefly.json index f2eb4400e0..709d557066 100644 --- a/schemas/pyrefly.json +++ b/schemas/pyrefly.json @@ -72,6 +72,12 @@ "type": "boolean", "default": true }, + "class-decorator-behavior": { + "description": "How class decorators affect class types. Ignore all decorator return types, ignore untyped decorators, or apply all decorator return types.", + "type": "string", + "enum": ["ignore", "typed-only", "all"], + "default": "typed-only" + }, "untyped-def-behavior": { "description": "How should Pyrefly treat function definitions with no parameter or return type annotations?", "type": "string", diff --git a/schemas/validate_schemas.py b/schemas/validate_schemas.py index 8f60a72788..e5bec27105 100644 --- a/schemas/validate_schemas.py +++ b/schemas/validate_schemas.py @@ -65,6 +65,21 @@ def _make_validator(schema_file: Path): class TestPositiveValidation(unittest.TestCase): """Positive tests: valid config files should pass schema validation.""" + def test_class_decorator_behavior(self) -> None: + validator = _make_validator(SCHEMAS_DIR / "pyrefly.json") + for value in ("ignore", "typed-only", "all"): + with self.subTest(value=value): + validator.validate({"class-decorator-behavior": value}) + validator.validate( + { + "sub-config": [ + {"matches": "legacy/**", "class-decorator-behavior": value} + ] + } + ) + with self.assertRaises(jsonschema.ValidationError): + validator.validate({"class-decorator-behavior": "invalid"}) + def test_pyrefly_toml(self) -> None: toml_file = SCHEMAS_DIR / "test-pyrefly.toml" schema_file = SCHEMAS_DIR / "pyrefly.json" diff --git a/website/docs/configuration.mdx b/website/docs/configuration.mdx index 6e506c3520..e1a9489cff 100644 --- a/website/docs/configuration.mdx +++ b/website/docs/configuration.mdx @@ -1114,6 +1114,40 @@ valid for a parameter, is ignored without an error. - Default: `false` - Flag equivalent: `--infer-param-types-from-docstring` +### `class-decorator-behavior` + +Controls how class decorators affect the type bound to a class name. + +- `"ignore"`: Ignore decorator return types and preserve the class type. Use this + mode for compatibility with mypy's treatment of class decorators. +- `"typed-only"` (default): Apply typed decorators. Preserve the class type + when a decorator has type `Any` or returns `Any` or `Unknown`. This mode follows + the class-decorator behavior used by Pyright and ty. An explicitly annotated + `type[Any]` return still erases the class interface. An inferred `type[Any]` + return preserves it. +- `"all"`: Apply all decorator return types, including `Any` and `Unknown`. + Decorators apply from the bottom of the decorator list to the top. + +For example, a decorator annotated to return `int` changes the class name's type +to `int` in `"typed-only"` and `"all"` modes. In `"ignore"` mode, the name +retains its class type. + +```toml +class-decorator-behavior = "typed-only" +``` + +Special decorators such as `dataclass`, `dataclass_transform`, `final`, +`disjoint_base`, `runtime_checkable`, and `total_ordering` retain their supported +class behavior in all three modes. An invalid decorator also preserves the class +during error recovery. + +This option changes inferred types. Configure the `untyped-class-decorator` +error kind separately to control diagnostics for untyped class decorators. + +- Type: one of `"ignore"`, `"typed-only"`, `"all"` +- Default: `"typed-only"` +- Flag equivalent: `--class-decorator-behavior` + ### `strict-callable-subtyping` Controls whether Pyrefly enforces strict parameter compatibility when checking @@ -1688,6 +1722,7 @@ We currently allow the following config options to be overridden in a SubConfig: - `replace-untyped-imports-with-any` - `check-unannotated-defs` - `infer-return-types` +- `class-decorator-behavior` - `ignore-errors-in-generated-code` Most SubConfig overrides _replace_ the values appearing in the 'root' or top-level of the From 0e9ba8f3a4bf85bb12f2dd54e9ff18416227905f Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 31 Oct 2026 05:41:54 +0900 Subject: [PATCH 5/6] fix generic --- pyrefly/lib/alt/function.rs | 22 ++++++----- pyrefly/lib/test/decorators.rs | 70 +++++++++++++++++++++++++++++++++- 2 files changed, 81 insertions(+), 11 deletions(-) diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index d2de37cd2d..452d96d988 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -2341,21 +2341,23 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } let call_target = self.as_call_target_or_error(decorator, CallStyle::FreeForm, range, errors, None); - let arg = CallArg::ty(&decoratee, range); + // 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; - // Retain class-definition metadata when the decorator returns the same nongeneric class. - // A specialized generic return must keep its inferred type arguments. - if let Type::ClassDef(cls) = &decoratee - && let Type::Type(inner) = &decorated - && let Type::ClassType(result) = &**inner - && result.class_object() == cls - && result.targs().is_empty() - { + if decorated == decoratee_arg { decoratee } else { - decorated + self.restore_decoratee_generics(decorated, decoratee_tparams) } } diff --git a/pyrefly/lib/test/decorators.rs b/pyrefly/lib/test/decorators.rs index d8f52103d5..e56f112d93 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -1273,6 +1273,16 @@ assert_type(Data(value=1).value, int) class Invalid: ... assert_type(Invalid(), Invalid) + +def generic_identity[T](cls: type[T]) -> type[T]: + return cls + +@generic_identity +class GenericClass[T]: + def method(self, value: T) -> T: + return value + +assert_type(GenericClass[int]().method(1), int) "#, ); @@ -1309,10 +1319,45 @@ assert_type(Child().inherited(), int) "#, ); +testcase!( + test_class_decorator_preserves_generic_class, + r#" +from typing import Callable, Generic, TypeVar, assert_type, overload + +T = TypeVar("T") + +# This is the overloaded class-preserving signature used by rich.repr.auto. +@overload +def auto(cls: type[T]) -> type[T]: ... +@overload +def auto(*, angular: bool = False) -> Callable[[type[T]], type[T]]: ... +def auto(cls=None, *, angular=False): + return cls if cls is not None else auto + +@auto +class Bare[T]: + def method(self, value: T) -> T: + return value + +@auto(angular=True) +class Factory(Generic[T]): + def method(self, value: T) -> T: + return value + +class BareChild(Bare[int]): ... +class FactoryChild(Factory[str]): ... + +assert_type(Bare[int]().method(1), int) +assert_type(Factory[str]().method("value"), str) +assert_type(BareChild().method(1), int) +assert_type(FactoryChild().method("value"), str) + "#, +); + testcase!( test_class_decorator_preserves_returned_type_arguments, r#" -from typing import Callable, assert_type, cast +from typing import Any, Callable, assert_type, cast def identity[T](cls: type[T]) -> type[T]: return cls @@ -1324,6 +1369,29 @@ class Box[T]: assert_type(Box(1).value, int) Box("bad") # E: is not assignable to parameter `value` with type `int` + +@cast("Callable[[type[Erased]], type[Erased[Any]]]", identity) +class Erased[T]: + def __init__(self, value: T) -> None: + self.value = value + +assert_type(Erased(1).value, Any) + "#, +); + +testcase!( + test_class_decorator_generic_noncallable_result, + r#" +from typing import Any, assert_type + +def instances[T](cls: type[T]) -> list[T]: ... + +@instances +class Box[T]: + value: T + +# Class type parameters cannot escape into a non-generic value. +assert_type(Box[0].value, Any) "#, ); From 8fdb554b5d9936866383516905ed86de4bac12af Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 6 Nov 2026 21:35:11 +0900 Subject: [PATCH 6/6] fix founded problems --- pyrefly/lib/alt/function.rs | 31 +++++++++++++++ pyrefly/lib/alt/solve.rs | 29 -------------- pyrefly/lib/test/decorators.rs | 70 ++++++++++++++++++++++++++++++++++ website/docs/configuration.mdx | 9 +++-- 4 files changed, 106 insertions(+), 33 deletions(-) diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index 452d96d988..19747c8e44 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -78,6 +78,7 @@ use crate::binding::binding::KeyDecorator; use crate::binding::binding::KeyLegacyTypeParam; use crate::binding::binding::KeyUndecoratedFunction; use crate::binding::classmethod::MethodReceiverKind; +use crate::config::base::ClassDecoratorBehavior; use crate::config::error_kind::ErrorKind; use crate::error::collector::ErrorCollector; use crate::error::context::TypeCheckContext; @@ -2310,6 +2311,36 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { if self.stdlib.is_bootstrapping() { return decoratee; } + if self.solver().config.class_decorator_behavior == ClassDecoratorBehavior::TypedOnly { + // Any untyped union member must preserve the class type. + let preserves_class = decorator.clone().into_unions().into_iter().any(|member| { + let instance_call = match &member { + Type::ClassType(instance) => self.instance_as_dunder_call(instance), + _ => None, + }; + let callable = instance_call.as_ref().unwrap_or(&member); + matches!(callable, Type::Any(_)) + || callable + .toplevel_callable_signatures() + .any(|(signature, _)| { + matches!(&signature.ret, Type::Any(_)) + || signature.ret.has_top_level_implicit_any() + || callable.toplevel_func_metadata().is_some_and(|meta| { + meta.flags.is_return_inferred + && matches!( + &signature.ret, + Type::Type(inner) if matches!( + &**inner, + Type::Any(AnyStyle::Explicit) + ) + ) + }) + }) + }); + if preserves_class { + return decoratee; + } + } let preserves_class = match &decorator { Type::KwCall(call) => { call.func_metadata diff --git a/pyrefly/lib/alt/solve.rs b/pyrefly/lib/alt/solve.rs index 97ae862145..d992d7bbe9 100644 --- a/pyrefly/lib/alt/solve.rs +++ b/pyrefly/lib/alt/solve.rs @@ -7532,35 +7532,6 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { if matches!(&decorator.ty, Type::Any(AnyStyle::Error)) { continue; } - if behavior == ClassDecoratorBehavior::TypedOnly { - let instance_call = match &decorator.ty { - Type::ClassType(instance) => self.instance_as_dunder_call(instance), - _ => None, - }; - let callable = instance_call.as_ref().unwrap_or(&decorator.ty); - let preserves_class = matches!(callable, Type::Any(_)) - || callable - .toplevel_callable_signatures() - .any(|(signature, _)| { - matches!(&signature.ret, Type::Any(_)) - || signature.ret.has_top_level_implicit_any() - || callable.toplevel_func_metadata().is_some_and( - |meta| { - meta.flags.is_return_inferred - && matches!( - &signature.ret, - Type::Type(inner) if matches!( - &**inner, - Type::Any(AnyStyle::Explicit) - ) - ) - }, - ) - }); - if preserves_class { - continue; - } - } let range = self.bindings().idx_to_key(*decorator_key).range(); ty = self.apply_class_decorator(decorator.ty.clone(), ty, range, errors); } diff --git a/pyrefly/lib/test/decorators.rs b/pyrefly/lib/test/decorators.rs index e56f112d93..2cb17c4b00 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -1197,6 +1197,76 @@ assert_type(Instance, type[Instance]) "#, ); +testcase!( + test_class_decorator_behavior_typed_only_union, + r#" +from typing import Any, Callable, assert_type + +def inferred_factory(value): + if value: + return lambda cls: cls + return value + +@inferred_factory(True) +class Inferred: + def method(self) -> int: + return 1 + +assert_type(Inferred, type[Inferred]) +assert_type(Inferred().method(), int) +value: str = Inferred().method() # E: `int` is not assignable to `str` + +def to_int(cls: type) -> int: ... +def to_str(cls: type) -> str: ... + +def typed_factory(flag: bool) -> Callable[[type], int] | Callable[[type], str]: + return to_int if flag else to_str + +@typed_factory(True) +class Typed: ... + +assert_type(Typed, int | str) + +def any_factory(value: Any) -> Callable[[type], int] | Any: + return value + +@any_factory(to_int) +class ExplicitAny: ... + +assert_type(ExplicitAny, type[ExplicitAny]) + +class UntypedDecorator: + def __call__(self, cls): + return cls + +def instance_factory(flag: bool) -> Callable[[type], int] | UntypedDecorator: + return to_int if flag else UntypedDecorator() + +@instance_factory(False) +class Instance: ... + +assert_type(Instance, type[Instance]) +"#, +); + +testcase!( + test_class_decorator_behavior_all_union, + TestEnv::new().with_class_decorator_behavior(ClassDecoratorBehavior::All), + r#" +from typing import Any, assert_type + +def factory(value): + if value: + return lambda cls: cls + return value + +@factory(True) +class C: ... + +assert_type(C, Any) +"#, +); + testcase!( test_class_decorator_behavior_default, r#" diff --git a/website/docs/configuration.mdx b/website/docs/configuration.mdx index e1a9489cff..7a3052b59c 100644 --- a/website/docs/configuration.mdx +++ b/website/docs/configuration.mdx @@ -1121,10 +1121,11 @@ Controls how class decorators affect the type bound to a class name. - `"ignore"`: Ignore decorator return types and preserve the class type. Use this mode for compatibility with mypy's treatment of class decorators. - `"typed-only"` (default): Apply typed decorators. Preserve the class type - when a decorator has type `Any` or returns `Any` or `Unknown`. This mode follows - the class-decorator behavior used by Pyright and ty. An explicitly annotated - `type[Any]` return still erases the class interface. An inferred `type[Any]` - return preserves it. + when a decorator has type `Any` or returns `Any` or `Unknown`. For a decorator + with a union type, any untyped member preserves the class type. This mode + follows the class-decorator behavior used by Pyright and ty. An explicitly + annotated `type[Any]` return still erases the class interface. An inferred + `type[Any]` return preserves it. - `"all"`: Apply all decorator return types, including `Any` and `Unknown`. Decorators apply from the bottom of the decorator list to the top.