feat: let ProxyMethod resolve callable attribute targets (#4915)

Summary:
Let `ProxyMethod` resolve a target declared as a class attribute, in addition to an ordinary instance method. A target is accepted whenever pyrefly can call it: `Callable[...]` annotations, callback protocols, and `type[X]` constructors all work but non-callable targets such as `forward: int` are still rejected.

This lets the shape stubs declare `nn.Module.forward` as `Callable[..., Any]`, matching PyTorch:

https://github.com/pytorch/pytorch/blob/d6c03540dc0f040e1d8a920d19d9333de6513545/torch/nn/modules/module.py#L526

The stubs previously declared `forward` as a method because that was the only target `ProxyMethod` could resolve. Under `strict-callable-subtyping`, a method-form `forward(*args, **kwargs)` makes every concrete `forward(self, x)` override a `bad-override`.

Pull Request resolved: #4915

Test Plan: Run `test.py`.

Reviewed By: stroxler

Differential Revision: D119805915

fbshipit-source-id: 5a7c941d69f3794e7896ee2bf0b3af9bd89e73a6

Author: yeetypete

pyrefly/lib/alt/attr.rs

@@ -286,7 +286,8 @@ pub enum NoAccessReason {
     SuperMethodNeedsImplementation(Class),
     /// A proxy method was accessed on the class object rather than an instance.
     ProxyMethodClassAccess(Class),
-    /// A proxy method declaration exists, but its target cannot be used as an instance method.
+    /// A proxy method declaration exists, but its target is neither an ordinary instance method
+    /// nor a class attribute whose type is callable.
     ProxyMethodTargetInvalid { class: Class, target: Name },
 }
 

pyrefly/lib/alt/class/class_field.rs

@@ -55,6 +55,7 @@ use crate::alt::answers_solver::AnswersSolver;
 use crate::alt::attr::AttrSubsetError;
 use crate::alt::attr::ClassBase;
 use crate::alt::attr::NoAccessReason;
+use crate::alt::call::CallTargetLookup;
 use crate::alt::callable::CallArg;
 use crate::alt::expr::TypeOrExpr;
 use crate::alt::types::class_bases::ClassBases;
@@ -2705,6 +2706,20 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         }
     }
 
+    /// A `ProxyMethod` target is either an ordinary instance method or a class
+    /// attribute whose type is callable (e.g. `torch.nn.Module` declares
+    /// `forward: Callable[..., Any]`).
+    fn is_proxy_method_target(&self, field: &ClassFieldInner) -> bool {
+        match field {
+            ClassFieldInner::Method { ty, .. } => Self::is_ordinary_instance_method_type(ty),
+            ClassFieldInner::ClassAttribute { ty, .. } => matches!(
+                self.as_call_target(self.normalize_attr_ty(ty.clone())),
+                CallTargetLookup::Ok(_)
+            ),
+            _ => false,
+        }
+    }
+
     fn is_ordinary_instance_method_type(ty: &Type) -> bool {
         ty.toplevel_func_metadata()
             .is_some_and(&|metadata: &FuncMetadata| {
@@ -3547,13 +3562,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             }
             ClassFieldInner::ProxyMethod { target, .. } => {
                 match self.get_class_member(instance.class, &target) {
-                    Some(target_field)
-                        if matches!(
-                            &target_field.0,
-                            ClassFieldInner::Method { ty, .. }
-                                if Self::is_ordinary_instance_method_type(ty)
-                        ) =>
-                    {
+                    Some(target_field) if self.is_proxy_method_target(&target_field.0) => {
                         self.as_instance_attribute(&target, target_field.as_ref(), instance)
                     }
                     _ => ClassAttribute::no_access(NoAccessReason::ProxyMethodTargetInvalid {

pyrefly/lib/test/protocol.rs

@@ -436,6 +436,62 @@ Base()()  # E: Proxy method `__call__` of class `Base` cannot resolve target met
 "#,
 );
 
+testcase!(
+    test_proxy_method_callable_attribute_target,
+    proxy_method_env(),
+    r#"
+from typing import Any, Callable, Protocol, assert_type
+from shape_extensions import ProxyMethod
+
+class CallableTarget:
+    __call__: ProxyMethod["forward"]
+    forward: Callable[[int], str]
+
+assert_type(CallableTarget()(1), str)
+CallableTarget()("bad")  # E: `Literal['bad']` is not assignable to parameter with type `int`
+
+class Callback(Protocol):
+    def __call__(self, x: int) -> str: ...
+
+class ProtocolTarget:
+    __call__: ProxyMethod["forward"]
+    forward: Callback
+
+assert_type(ProtocolTarget()(1), str)
+
+class GradualCallableTarget:
+    __call__: ProxyMethod["forward"]
+    forward: Callable[..., Any]
+
+assert_type(GradualCallableTarget()(1, "two", three=3), Any)
+"#,
+);
+
+testcase!(
+    test_proxy_method_callable_attribute_target_strict_override,
+    proxy_method_env().enable_strict_callable_subtyping(),
+    r#"
+from typing import Any, Callable, assert_type
+from shape_extensions import ProxyMethod
+
+class AttributeBase:
+    __call__: ProxyMethod["forward"]
+    forward: Callable[..., Any]
+
+class AttributeChild(AttributeBase):
+    def forward(self, x: int) -> str: ...
+
+assert_type(AttributeChild()(1), str)
+
+class MethodBase:
+    __call__: ProxyMethod["forward"]
+    def forward(self, *args: Any, **kwargs: Any) -> Any: ...
+
+class MethodChild(MethodBase):
+    def forward(self, x: int) -> str: ...  # E: overrides parent class `MethodBase` in an inconsistent manner
+"#,
+);
+
 testcase!(
     test_proxy_method_rejects_proxy_chain_and_self_reference,
     proxy_method_env(),

tensor-shapes/pyrefly-torch-stubs/suites.py

@@ -20,7 +20,11 @@
 # torch runtime tests are separate unittest modules under test/runtime_tests.
 SUITES: list[Suite] = [
     Suite(name="torch-examples", patterns=("examples/*.py", "examples/runtime/*.py")),
-    Suite(name="torch-positive", patterns=("test/test_*.py",)),
+    Suite(
+        name="torch-positive",
+        patterns=("test/test_*.py",),
+        strict_callable_subtyping=True,
+    ),
     Suite(
         name="torch-negative",
         patterns=("test/negative_tests/test_*.py",),

tensor-shapes/pyrefly-torch-stubs/test/test_module_forward_attribute.py

@@ -0,0 +1,51 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Test `nn.Module.forward` declared as a callable attribute."""
+
+from typing import Any, assert_type, override, TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from shape_extensions import Int, IntVar
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+
+class LinearLayer[N: IntVar, M: IntVar](nn.Module):
+    def __init__(self, n: Int[N], m: Int[M]) -> None:
+        super().__init__()
+        self.linear = nn.Linear(n, m)
+
+    @override
+    def forward[B: IntVar](self, x: Tensor[[B, N]]) -> Tensor[[B, M]]:
+        return self.linear(x)
+
+
+class Passthrough(nn.Module):
+    """Transform-style module, like torchvision's `Transform.forward(self, *inputs)`."""
+
+    @override
+    def forward(self, *inputs: Any) -> Any:
+        return inputs
+
+
+def test_forward_override_keeps_call_proxy() -> None:
+    x: Tensor[[16, 6]] = torch.randn(16, 6)
+    layer = LinearLayer(6, 9)
+    assert_type(layer(x), Tensor[[16, 9]])
+    assert_type(layer.forward(x), Tensor[[16, 9]])
+
+
+def test_call_through_base_module_type() -> None:
+    x: Tensor[[16, 6]] = torch.randn(16, 6)
+    module: nn.Module = LinearLayer(6, 9)
+    assert_type(module(x), Any)
+
+
+def test_forward_override_with_star_args() -> None:
+    module = Passthrough()
+    module(torch.zeros(2), torch.zeros(3))

tensor-shapes/pyrefly-torch-stubs/torch-stubs/nn/__init__.pyi

@@ -157,7 +157,7 @@ class Module:
     def __getattr__(self, name: str) -> Any: ...
     def __setattr__(self, name: str, value: Any) -> None: ...
     __call__: ProxyMethod["forward"]
-    def forward(self, *args: Any, **kwargs: Any) -> Any: ...
+    forward: Callable[..., Any]
     def extra_repr(self) -> str: ...
     def register_buffer(
         self, name: str, tensor: Tensor | None, persistent: bool = True

tensor-shapes/shape_testing.py

@@ -69,6 +69,7 @@ class Suite:
     # are otherwise suppressed, which is why the torch corpus enables it only
     # for the dedicated negative-test directories.
     expectations: bool = False
+    strict_callable_subtyping: bool = False
 
     def files(self, package_root: Path) -> list[str]:
         paths = sorted(
@@ -252,6 +253,8 @@ def check_suites(
         ]
         if suite.expectations:
             command.append("--expectations")
+        if suite.strict_callable_subtyping:
+            command.append("--strict-callable-subtyping=true")
         for search_path in (
             *suite.extra_search_paths,
             package_root,