From c7f3fa72cedcde682c87145272d82f1d34ebec34 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 30 Oct 2026 19:24:32 +0900 Subject: [PATCH 1/2] test --- pyrefly/lib/test/decorators.rs | 136 +++++++++++++++++++++++++++++++-- 1 file changed, 128 insertions(+), 8 deletions(-) diff --git a/pyrefly/lib/test/decorators.rs b/pyrefly/lib/test/decorators.rs index 99904106c8..9a0ad2e7b8 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -781,22 +781,21 @@ class C(ABC): // These assignments should behave like @classmethod: cls should be type[Self], // calls should preserve subclass types, and initialized attributes should exist on the class. testcase!( - bug = "Desugared classmethod assignments retain instance binding", test_desugared_classmethod_application, r#" from typing import Self, assert_type class A: def f(cls): - assert_type(cls, type[Self]) # E: assert_type(Self@A, type[Self@A]) failed + assert_type(cls, type[Self]) return cls f = classmethod(f) class B(A): pass -assert_type(A.f(), type[A]) # E: assert_type(A, type[A]) failed # E: Argument `type[A]` is not assignable to parameter `cls` with type `A` in function `A.f` -assert_type(A().f(), type[A]) # E: assert_type(A, type[A]) failed # E: Argument `type[A]` is not assignable to parameter `cls` with type `A` in function `A.f` -assert_type(B.f(), type[B]) # E: assert_type(B, type[B]) failed # E: Argument `type[B]` is not assignable to parameter `cls` with type `B` in function `A.f` +assert_type(A.f(), type[A]) +assert_type(A().f(), type[A]) +assert_type(B.f(), type[B]) class Holder: def foo(cls, string: str) -> int: @@ -804,7 +803,7 @@ class Holder: foo = classmethod(foo) -assert_type(Holder.foo("hello"), int) # E: Argument `type[Holder]` is not assignable to parameter `cls` with type `Holder` in function `Holder.foo` +assert_type(Holder.foo("hello"), int) class InitializesOnClass: def initialize(cls) -> None: @@ -812,8 +811,129 @@ class InitializesOnClass: initialize = classmethod(initialize) -InitializesOnClass.initialize() # E: Argument `type[InitializesOnClass]` is not assignable to parameter `cls` with type `InitializesOnClass` in function `InitializesOnClass.initialize` -assert_type(InitializesOnClass.value, int) # E: assert_type(Unknown, int) failed # E: Instance-only attribute `value` of class `InitializesOnClass` is not visible on the class +InitializesOnClass.initialize() +assert_type(InitializesOnClass.value, int) + "#, +); + +testcase!( + test_desugared_classmethod_scope, + r#" +from typing import assert_type + +def f(cls: type[object]) -> int: + return 1 + +class Outer: + def f(self): + return self + + class Inner: + f = classmethod(f) + +assert_type(Outer().f(), Outer) +assert_type(Outer.Inner.f(), int) + "#, +); + +testcase!( + test_desugared_classmethod_alias, + r#" +from typing import Self, assert_type + +class A: + def f(self): + return self + + g = f + f = classmethod(f) + +class B(A): + pass + +assert_type(A().g(), A) +assert_type(B().g(), B) +assert_type(A.f(), type[A]) +assert_type(B.f(), type[B]) + +class Redefined: + def f(self): + return self + g = f + + def f(cls): + return cls + f = classmethod(f) + +assert_type(Redefined().g(), Redefined) +assert_type(Redefined.f(), type[Redefined]) + +class Generic[T = str]: + def f(self): + return self + + g = f + f = classmethod(f) + +assert_type(Generic[int]().g(), Generic[int]) +assert_type(Generic[int].f(), type[Generic[int]]) + +class WithSelf: + def f(self, other: Self) -> Self: + return other + + g = f + f = classmethod(f) + +assert_type(WithSelf().g(WithSelf()), WithSelf) +assert_type(WithSelf.f(WithSelf()), WithSelf) + "#, +); + +testcase!( + test_desugared_classmethod_decorator_composition, + r#" +from typing import Callable, Concatenate, assert_type + +def adapt[**P, R](f: Callable[Concatenate[A, P], R]) -> Callable[Concatenate[type[A], P], R]: + def wrapper(cls: type[A], *args: P.args, **kwargs: P.kwargs) -> R: + return f(cls(), *args, **kwargs) + return wrapper + +class A: + @adapt + def f(self, x: int) -> int: + self.value = x + return x + + f = classmethod(f) + +assert_type(A.f(1), int) +assert_type(A().f(1), int) +A.value # E: Instance-only attribute `value` of class `A` is not visible on the class + "#, +); + +testcase!( + test_desugared_classmethod_overload, + r#" +from typing import assert_type, overload + +class A: + @overload + def f(cls, x: int) -> int: ... + @overload + def f(cls, x: str) -> str: ... + def f(cls, x: int | str) -> int | str: + return x + + g = f + f = classmethod(f) + +assert_type(A.f(1), int) +assert_type(A().f("hello"), str) +assert_type(A().g(1), int) +assert_type(A().g("hello"), str) "#, ); From 711f25566a0f4ebbb92002ee77c210eadd33d8d5 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Thu, 3 Sep 2026 15:02:38 +0900 Subject: [PATCH 2/2] fix --- pyrefly/lib/alt/function.rs | 47 ++++++++++++++++++++++- pyrefly/lib/binding/binding.rs | 6 +++ pyrefly/lib/binding/bindings.rs | 66 ++++++++++++++++++++++++++++++++- pyrefly/lib/binding/expr.rs | 15 ++++++++ pyrefly/lib/binding/scope.rs | 16 ++++++++ pyrefly/lib/binding/stmt.rs | 8 ++++ 6 files changed, 156 insertions(+), 2 deletions(-) diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index 3d91c68246..ee9110e04d 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -29,9 +29,12 @@ use pyrefly_types::literal::LitStyle; use pyrefly_types::meta_shape_dsl::ShapeDslFunction; use pyrefly_types::meta_shape_dsl::ShapeTransform; use pyrefly_types::meta_shape_dsl::validate_shape_dsl_functions; +use pyrefly_types::quantified::AnchorIndex; use pyrefly_types::quantified::Quantified; +use pyrefly_types::quantified::QuantifiedIdentity; use pyrefly_types::quantified::QuantifiedOrigin; use pyrefly_types::type_level_dsl::ParsedTypeShapeDslFunction; +use pyrefly_types::type_var::PreInferenceVariance; use pyrefly_types::type_var::Restriction; use pyrefly_types::types::AnyStyle; use pyrefly_types::types::BoundMethod; @@ -600,7 +603,43 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // Look for a @classmethod or @staticmethod decorator and change the "self" type // accordingly. This is not totally correct, since it doesn't account for chaining // decorators, or weird cases like both decorators existing at the same time. - if flags.is_classmethod || found_class_property || is_dunder_new { + let mut receiver_tparam = None; + if def.is_reassigned_classmethod + && def.is_used_before_classmethod + && def + .parameters + .iter() + .next() + .is_some_and(|p| !p.is_variadic() && p.annotation().is_none()) + { + if let Some(cls) = &defining_cls { + // An earlier alias can bind an instance while classmethod binds a class. + // Quantify the receiver so each call retains its own return type. + let mut instance = self.heap.mk_class_type(self.as_class_type_unchecked(cls)); + // The bound accepts every specialization of the defining class. + instance.subst_mut_fn(&mut |q| { + Some(Quantified::as_gradual_type_helper(q.kind(), None)) + }); + let class = self.heap.mk_type_of(instance.clone()); + let receiver = Quantified::type_var( + Name::new_static("Receiver"), + QuantifiedIdentity::new( + self.module().name(), + AnchorIndex::first(def.name.range), + QuantifiedOrigin::Synthetic { is_self: true }, + ), + None, + Restriction::Bound(self.unions(vec![instance, class])), + PreInferenceVariance::Undefined, + ); + self_type = Some(self.heap.mk_quantified(receiver.clone())); + receiver_tparam = Some(receiver); + } + } else if flags.is_classmethod + || found_class_property + || is_dunder_new + || def.is_reassigned_classmethod + { self_type = self_type.map(|t| self.heap.mk_type_of(t)); } else if flags.is_staticmethod { self_type = None; @@ -698,6 +737,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { .iter() .filter_map(|key| self.get_idx(*key).parameter().cloned()); tparams.extend(legacy_tparams); + tparams.extend(receiver_tparam); // Validate written parameters before declaration-scoped dimensions are appended below. let tparams = Arc::new(self.validated_tparams(def.range, tparams, TParamsSource::Function, errors)); @@ -891,6 +931,11 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // For classmethods, the annotation is `type[TFoo]`, so we also unwrap `Type::Type(...)`. if has_implicit_self_or_cls_param && !def.metadata.flags.is_staticmethod + && stmt + .parameters + .iter() + .next() + .is_some_and(|p| p.annotation().is_some()) && let Some(first_param) = def.params.first() && { let ty = first_param.as_type(); diff --git a/pyrefly/lib/binding/binding.rs b/pyrefly/lib/binding/binding.rs index fe4e29bbce..d97c6938cd 100644 --- a/pyrefly/lib/binding/binding.rs +++ b/pyrefly/lib/binding/binding.rs @@ -2076,6 +2076,10 @@ pub struct FunctionDefData { pub parameters: Box, pub type_params: Option>, pub is_async: bool, + /// Whether this method is converted with the class-body idiom `f = classmethod(f)`. + pub is_reassigned_classmethod: bool, + /// Whether another expression uses this function before its classmethod assignment. + pub is_used_before_classmethod: bool, pub range: TextRange, } @@ -2086,6 +2090,8 @@ impl FunctionDefData { parameters: def.parameters, type_params: def.type_params, is_async: def.is_async, + is_reassigned_classmethod: false, + is_used_before_classmethod: false, range: def.range, } } diff --git a/pyrefly/lib/binding/bindings.rs b/pyrefly/lib/binding/bindings.rs index cfe37b1610..483e7b2202 100644 --- a/pyrefly/lib/binding/bindings.rs +++ b/pyrefly/lib/binding/bindings.rs @@ -1134,7 +1134,7 @@ impl<'a> BindingsBuilder<'a> { matches!(self.idx_to_binding(idx), Some(Binding::ClassDef(..))) } - fn idx_to_binding_mut(&mut self, idx: Idx) -> Option<&mut K::Value> + pub(super) fn idx_to_binding_mut(&mut self, idx: Idx) -> Option<&mut K::Value> where K: Keyed, BindingTable: TableKeyed>, @@ -1142,6 +1142,70 @@ impl<'a> BindingsBuilder<'a> { self.table.get_mut::().1.get_mut(idx) } + /// Infer the receiver for a local classmethod assignment without changing method flags. + pub fn infer_classmethod_receiver(&mut self, name: &Name) { + let Some(mut idx) = self.scopes.current_flow_idx(name) else { + return; + }; + let mut functions = Vec::new(); + while let Some(Binding::Function { + decorated_idx, + pred_idx, + .. + }) = self.idx_to_binding(idx) + { + let pred_idx = *pred_idx; + let undecorated_idx = self + .idx_to_binding(*decorated_idx) + .expect("a function flow binding has a decorated function") + .undecorated_idx; + let binding = self + .idx_to_binding(undecorated_idx) + .expect("a decorated function has an undecorated function"); + let specials = binding + .decorators + .iter() + .map(|idx| { + let decorator = self + .idx_to_binding(*idx) + .expect("a decorator has a binding"); + self.as_special_export(&decorator.expr) + }) + .collect::>(); + if !functions.is_empty() && !specials.contains(&Some(SpecialExport::Overload)) { + break; + } + // Custom decorators receive the original function and can change its receiver. + if specials + .iter() + .any(|special| *special != Some(SpecialExport::Overload)) + { + return; + } + functions.push(undecorated_idx); + let Some(pred_idx) = pred_idx else { break }; + idx = pred_idx; + } + if functions.is_empty() { + return; + } + let is_used = functions.iter().any(|idx| { + self.idx_to_binding(*idx) + .expect("a decorated function has an undecorated function") + .def + .is_used_before_classmethod + }); + for idx in functions { + let function = &mut self + .idx_to_binding_mut(idx) + .expect("a decorated function has an undecorated function") + .def; + function.is_reassigned_classmethod = true; + function.is_used_before_classmethod = is_used; + } + self.scopes.mark_method_as_classmethod(name); + } + /// Declare a `Key` as a usage, which can be used for name lookups. Like `idx_for_promise`, /// this is a promise to later provide a `Binding` corresponding this key. pub fn declare_current_idx(&mut self, key: Key) -> CurrentIdx { diff --git a/pyrefly/lib/binding/expr.rs b/pyrefly/lib/binding/expr.rs index 65e04242ea..352b4b20c8 100644 --- a/pyrefly/lib/binding/expr.rs +++ b/pyrefly/lib/binding/expr.rs @@ -420,6 +420,21 @@ impl<'a> BindingsBuilder<'a> { is_module_scope, is_outer_class_type_parameter, } => { + if self.scopes.in_class_body() + && let Some(Binding::Function { decorated_idx, .. }) = + self.idx_to_binding(value_idx.unwrap_or(lookup_result_idx)) + { + let undecorated_idx = self + .idx_to_binding(*decorated_idx) + .expect("a function flow binding has a decorated function") + .undecorated_idx; + let function = self + .idx_to_binding_mut(undecorated_idx) + .expect("a decorated function has an undecorated function"); + if !function.def.is_reassigned_classmethod { + function.def.is_used_before_classmethod = true; + } + } if is_outer_class_type_parameter { return self.insert_binding( key, diff --git a/pyrefly/lib/binding/scope.rs b/pyrefly/lib/binding/scope.rs index bbbb437a3e..49424ce239 100644 --- a/pyrefly/lib/binding/scope.rs +++ b/pyrefly/lib/binding/scope.rs @@ -1792,6 +1792,22 @@ impl Scopes { } } + pub fn mark_method_as_classmethod(&mut self, method_name: &Name) { + let ScopeKind::Class(class_scope) = &mut self.current_mut().kind else { + unreachable!("a reassigned classmethod is only marked in a class body") + }; + for attributes in [ + &mut class_scope.attributes_from_recognized_methods, + &mut class_scope.attributes_from_other_methods, + ] { + if let Some(attributes) = attributes.get_mut(method_name) { + for attribute in attributes.values_mut() { + attribute.3 = MethodSelfKind::Class; + } + } + } + } + /// The `ClassDefIndex` of the current class body, if the innermost scope is one. pub fn current_class_def_index(&self) -> Option { match &self.current().kind { diff --git a/pyrefly/lib/binding/stmt.rs b/pyrefly/lib/binding/stmt.rs index 3f29b59025..0a25422ba4 100644 --- a/pyrefly/lib/binding/stmt.rs +++ b/pyrefly/lib/binding/stmt.rs @@ -1020,6 +1020,14 @@ impl<'a> BindingsBuilder<'a> { self.assign_type_var(name, call, kind); return; } + if special == SpecialExport::ClassMethod + && self.scopes.in_class_body() + && call.arguments.keywords.is_empty() + && let [Expr::Name(argument)] = &*call.arguments.args + && argument.id == name.id + { + self.infer_classmethod_receiver(&name.id); + } match special { SpecialExport::ParamSpec => { self.assign_param_spec(name, call);