_restricted_classmethod: add wrapper, to allow inspection (#19332)

Author: lauritsf

src/lightning/pytorch/CHANGELOG.md

@@ -77,6 +77,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed an issue with the ModelCheckpoint callback not saving relative symlinks with `ModelCheckpoint(save_last="link")` ([#19303](https://github.com/Lightning-AI/lightning/pull/19303))
 
 
+- Fixed issue where the `_restricted_classmethod_impl` would incorrectly raise a TypeError on inspection rather than on call ([#19332](https://github.com/Lightning-AI/lightning/pull/19332))
+
+
 ## [2.1.3] - 2023-12-21
 
 ### Changed

src/lightning/pytorch/utilities/model_helpers.py

@@ -11,10 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import functools
 import inspect
 import logging
 import os
-from types import MethodType
 from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Optional, Type, TypeVar
 
 from lightning_utilities.core.imports import RequirementCache
@@ -108,18 +108,23 @@ class _restricted_classmethod_impl(Generic[_T, _P, _R_co]):
     """Drop-in replacement for @classmethod, but raises an exception when the decorated method is called on an instance
     instead of a class type."""
 
-    def __init__(self, method: Callable[Concatenate[_T, _P], _R_co]) -> None:
+    def __init__(self, method: Callable[Concatenate[Type[_T], _P], _R_co]) -> None:
         self.method = method
 
     def __get__(self, instance: Optional[_T], cls: Type[_T]) -> Callable[_P, _R_co]:
-        # Workaround for https://github.com/pytorch/pytorch/issues/67146
-        is_scripting = any(os.path.join("torch", "jit") in frameinfo.filename for frameinfo in inspect.stack())
-        if instance is not None and not is_scripting:
-            raise TypeError(
-                f"The classmethod `{cls.__name__}.{self.method.__name__}` cannot be called on an instance."
-                " Please call it on the class type and make sure the return value is used."
-            )
-        return MethodType(self.method, cls)
+        # The wrapper ensures that the method can be inspected, but not called on an instance
+        @functools.wraps(self.method)
+        def wrapper(*args: Any, **kwargs: Any) -> _R_co:
+            # Workaround for https://github.com/pytorch/pytorch/issues/67146
+            is_scripting = any(os.path.join("torch", "jit") in frameinfo.filename for frameinfo in inspect.stack())
+            if instance is not None and not is_scripting:
+                raise TypeError(
+                    f"The classmethod `{cls.__name__}.{self.method.__name__}` cannot be called on an instance."
+                    " Please call it on the class type and make sure the return value is used."
+                )
+            return self.method(cls, *args, **kwargs)
+
+        return wrapper
 
 
 # trick static type checkers into thinking it's a @classmethod

tests/tests_pytorch/utilities/test_model_helpers.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import inspect
 import logging
 
 import pytest
@@ -66,10 +67,12 @@ def cmethod(cls):
 
 
 def test_restricted_classmethod():
+    restricted_method = RestrictedClass().restricted_cmethod  # no exception when getting restricted method
+
     with pytest.raises(TypeError, match="cannot be called on an instance"):
-        RestrictedClass().restricted_cmethod()
+        restricted_method()
 
-    RestrictedClass.restricted_cmethod()  # no exception
+    _ = inspect.getmembers(RestrictedClass())  # no exception on inspecting instance
 
 
 def test_module_mode():