fix foo = classmethod(foo) does not create a working classmethod #3362

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();

pyrefly/lib/binding/binding.rs

@@ -2076,6 +2076,10 @@ pub struct FunctionDefData {
     pub parameters: Box<Parameters>,
     pub type_params: Option<Box<TypeParams>>,
     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,
         }
     }

pyrefly/lib/binding/bindings.rs

@@ -1134,14 +1134,78 @@ impl<'a> BindingsBuilder<'a> {
         matches!(self.idx_to_binding(idx), Some(Binding::ClassDef(..)))
     }
 
-    fn idx_to_binding_mut<K>(&mut self, idx: Idx<K>) -> Option<&mut K::Value>
+    pub(super) fn idx_to_binding_mut<K>(&mut self, idx: Idx<K>) -> Option<&mut K::Value>
     where
         K: Keyed,
         BindingTable: TableKeyed<K, Value = BindingEntry<K>>,
     {
         self.table.get_mut::<K>().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::<Vec<_>>();
+            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 {

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,

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<ClassDefIndex> {
         match &self.current().kind {

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);