diff --git a/crates/pyrefly_types/src/quantified.rs b/crates/pyrefly_types/src/quantified.rs index 98eb72a3e2..0cff09825e 100644 --- a/crates/pyrefly_types/src/quantified.rs +++ b/crates/pyrefly_types/src/quantified.rs @@ -44,6 +44,8 @@ pub enum QuantifiedOrigin { /// Is this a Self quantified synthesized for `__new__` on a class? is_self: bool, }, + /// An inferred receiver shared by an instance method and a classmethod wrapper. + MethodReceiver, /// Synthetic binder for the parameter of an `IntTuples` mapper. MapIntTuplesParameter, /// De Bruijn sentinel used only while comparing `MapIntTuples` lambdas. diff --git a/pyrefly/lib/alt/attr.rs b/pyrefly/lib/alt/attr.rs index 5a4c6782e2..2fa299857d 100644 --- a/pyrefly/lib/alt/attr.rs +++ b/pyrefly/lib/alt/attr.rs @@ -59,6 +59,7 @@ use crate::types::literal::Lit; use crate::types::module::ModuleType; use crate::types::quantified::Quantified; use crate::types::quantified::QuantifiedKind; +use crate::types::quantified::QuantifiedOrigin; use crate::types::read_only::ReadOnlyReason; use crate::types::type_info::TypeInfo; use crate::types::type_var::Restriction; @@ -2548,6 +2549,21 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { quantified: Quantified, bound: AttributeBase1, ) -> Option { + if quantified.identity().origin == QuantifiedOrigin::MethodReceiver { + // Preserve each bound's descriptor binding behavior. Use Self so + // classmethod results refer to the class, not type[Receiver]. + return Some(match bound { + AttributeBase1::ClassInstance(cls) => { + AttributeBase1::SelfType(self.as_class_type_unchecked(cls.class_object())) + } + AttributeBase1::ClassObject(cls) | AttributeBase1::GenericAlias(cls) => { + AttributeBase1::ClassObject(ClassBase::SelfType( + self.as_class_type_unchecked(cls.class_object()), + )) + } + _ => unreachable!("a method receiver is bounded by its class and instances"), + }); + } match bound { AttributeBase1::Any(style) => Some(AttributeBase1::Any(style)), AttributeBase1::ClassObject(ClassBase::ClassDef(cls) | ClassBase::ClassType(cls)) => { diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index 38cd7bd1bb..8ee96e0ca2 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::MethodReceiver, + ), + 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,9 +737,13 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { .iter() .filter_map(|key| self.get_idx(*key).parameter().cloned()); tparams.extend(legacy_tparams); - // Validate written parameters before declaration-scoped dimensions are appended below. - let tparams = - Arc::new(self.validated_tparams(def.range, tparams, TParamsSource::Function, errors)); + // Only written parameters are subject to default ordering constraints. + let mut tparams = + self.validated_tparams(def.range, tparams, TParamsSource::Function, errors); + if let Some(receiver) = receiver_tparam { + tparams.extend(&TParams::new(vec![receiver])); + } + let tparams = Arc::new(tparams); let func_id = Arc::new(FuncDefId { qname: QName::new(def.name.clone(), parent.dupe(), self.module().dupe()), cls: defining_cls.clone(), @@ -891,6 +934,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 33c1e032cb..7b422f915f 100644 --- a/pyrefly/lib/binding/binding.rs +++ b/pyrefly/lib/binding/binding.rs @@ -2081,6 +2081,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, } @@ -2091,6 +2095,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 a39712e96a..4ab84bbd6a 100644 --- a/pyrefly/lib/binding/bindings.rs +++ b/pyrefly/lib/binding/bindings.rs @@ -1189,7 +1189,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>, @@ -1197,6 +1197,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 9b06754632..0c8d02b0ea 100644 --- a/pyrefly/lib/binding/expr.rs +++ b/pyrefly/lib/binding/expr.rs @@ -420,6 +420,23 @@ impl<'a> BindingsBuilder<'a> { is_module_scope, is_outer_class_type_parameter, } => { + if let Some(Binding::Function { + decorated_idx, + in_class: true, + .. + }) = 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 084cf0a01c..1947534808 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); diff --git a/pyrefly/lib/test/decorators.rs b/pyrefly/lib/test/decorators.rs index 8c6d0cc631..8fe6745621 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -1284,22 +1284,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: @@ -1307,7 +1306,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: @@ -1315,8 +1314,235 @@ 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_captured_default, + r#" +from typing import Callable, assert_type + +class A: + def f(self) -> int: + return 1 + + def use(self, callback: Callable[[A], int] = f) -> int: + return callback(self) + + f = classmethod(f) + +assert_type(A().use(), int) +assert_type(A.f(), int) + "#, +); + +testcase!( + test_desugared_classmethod_receiver_dependent_return, + r#" +from typing import assert_type + +class A: + @classmethod + def helper(cls): + return cls + + def f(self): + return self.helper() + + g = f + f = classmethod(f) + +class B(A): + pass + +assert_type(A.f(), type[A]) +assert_type(A().g(), type[A]) +assert_type(B.f(), type[B]) +assert_type(B().g(), type[B]) + +class Generic[T]: + @classmethod + def helper(cls): + return cls + + def f(self): + return self.helper() + + g = f + f = classmethod(f) + +assert_type(Generic[int].f(), type[Generic[int]]) +assert_type(Generic[int]().g(), type[Generic[int]]) + +class Specialized(Generic[int]): + pass + +assert_type(Specialized.f(), type[Specialized]) +assert_type(Specialized().g(), type[Specialized]) + +class Factory[T]: + @classmethod + def helper(cls): + return cls() + + def f(self): + return self.helper() + + g = f + f = classmethod(f) + +assert_type(Factory[str].f(), Factory[str]) +assert_type(Factory[str]().g(), Factory[str]) + "#, +); + +testcase!( + test_desugared_classmethod_defaulted_type_parameter, + r#" +from typing import assert_type + +class A: + def f[T = int](self, x: T) -> T: + return x + + g = f + f = classmethod(f) + +assert_type(A.f(1), int) +assert_type(A().g("hello"), str) + +class SameName: + def f[Receiver = int](self, x: Receiver) -> Receiver: + return x + + g = f + f = classmethod(f) + +assert_type(SameName.f(1), int) +assert_type(SameName().g("hello"), str) + "#, +); + +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) "#, );