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

pyrefly/lib/alt/function.rs

@@ -480,6 +480,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         body_kind: BodyKind,
         is_return_inferred: bool,
         calls_super_method: bool,
+        is_reassigned_classmethod: bool,
         class_key: Option<&Idx<KeyClass>>,
         decorators: &[Idx<KeyDecorator>],
         legacy_tparams: &[Idx<KeyLegacyTypeParam>],
@@ -504,7 +505,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
 
         let mut flags = FuncFlags {
             is_staticmethod: is_dunder_new,
-            is_classmethod: is_dunder_init_subclass,
+            is_classmethod: is_dunder_init_subclass || is_reassigned_classmethod,
             is_async: def.is_async,
             body_kind,
             is_return_inferred,

pyrefly/lib/alt/solve.rs

@@ -6273,6 +6273,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             x.body_kind,
             x.is_return_inferred,
             x.calls_super_method,
+            x.def.is_reassigned_classmethod,
             x.class_key.as_ref(),
             &x.decorators,
             &x.legacy_tparams,

pyrefly/lib/binding/binding.rs

@@ -1900,6 +1900,8 @@ 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,
     pub range: TextRange,
 }
 
@@ -1910,6 +1912,7 @@ impl FunctionDefData {
             parameters: def.parameters,
             type_params: def.type_params,
             is_async: def.is_async,
+            is_reassigned_classmethod: false,
             range: def.range,
         }
     }

pyrefly/lib/binding/bindings.rs

@@ -1026,6 +1026,17 @@ impl<'a> BindingsBuilder<'a> {
         self.table.get_mut::<K>().1.get_mut(idx)
     }
 
+    pub fn mark_reassigned_classmethod(&mut self, function_idx: Idx<KeyDecoratedFunction>) {
+        let undecorated_idx = self
+            .idx_to_binding(function_idx)
+            .expect("a function flow binding has a decorated function")
+            .undecorated_idx;
+        self.idx_to_binding_mut(undecorated_idx)
+            .expect("a decorated function has an undecorated function")
+            .def
+            .is_reassigned_classmethod = true;
+    }
+
     /// 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/scope.rs

@@ -1075,6 +1075,19 @@ impl ScopeClass {
         }
     }
 
+    fn mark_method_as_classmethod(&mut self, method_name: &Name) {
+        for attributes in [
+            &mut self.attributes_from_recognized_methods,
+            &mut self.attributes_from_other_methods,
+        ] {
+            if let Some(attributes) = attributes.get_mut(method_name) {
+                for attribute in attributes.values_mut() {
+                    attribute.3 = MethodSelfKind::Class;
+                }
+            }
+        }
+    }
+
     /// Produces triples (hashed_attr_name, MethodThatSetsAttr, attribute) for all assignments
     /// to `self.<attr_name>` in methods.
     ///
@@ -1718,6 +1731,13 @@ 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")
+        };
+        class_scope.mark_method_as_classmethod(method_name);
+    }
+
     /// 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

@@ -759,6 +759,17 @@ 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
+                            && let Some((_, FlowStyle::FunctionDef { function_idx, .. })) =
+                                self.scopes.binding_idx_for_name(&name.id)
+                        {
+                            self.mark_reassigned_classmethod(function_idx);
+                            self.scopes.mark_method_as_classmethod(&name.id);
+                        }
                         match special {
                             SpecialExport::ParamSpec => {
                                 self.assign_param_spec(name, call);

pyrefly/lib/test/decorators.rs

@@ -667,16 +667,61 @@ class C(ABC):
     "#,
 );
 
+// Regression test for https://github.com/facebook/pyrefly/issues/3362.
 testcase!(
-    bug = "We should treat `A.f` as a classmethod",
-    test_desugared_decorator_application,
+    test_desugared_classmethod_application,
     r#"
-from typing import assert_type
+from typing import Self, assert_type
 class A:
     def f(cls):
+        assert_type(cls, type[Self])
         return cls
     f = classmethod(f)
-assert_type(A.f(), type[A])  # E: assert_type(A, type[A])  # E: `type[A]` is not assignable to parameter `cls` with type `A`
+
+class B(A):
+    pass
+
+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:
+        return len(string)
+
+    foo = classmethod(foo)
+
+assert_type(Holder.foo("hello"), int)
+
+class InitializesOnClass:
+    def initialize(cls) -> None:
+        cls.value = 1
+
+    initialize = classmethod(initialize)
+
+InitializesOnClass.initialize()
+assert_type(InitializesOnClass.value, int)
+    "#,
+);
+
+testcase!(
+    test_desugared_classmethod_respects_shadowing,
+    r#"
+from typing import Self, assert_type
+
+def identity[T](x: T) -> T:
+    return x
+
+class A:
+    classmethod = identity
+
+    def f(self):
+        assert_type(self, Self)
+        return self
+
+    f = classmethod(f)
+
+assert_type(A().f(), A)
     "#,
 );