Fixed Broken MethodMonitor

Author: HMohammad2520

classmods/_method_monitor.py

@@ -1,20 +1,27 @@
-from typing import Any, Dict, List, Tuple, Type, Callable
+from typing import Any, Dict, List, Tuple, Type, Callable, Optional
 from functools import wraps
+import inspect
 
 class MethodMonitor:
-    # Dictionary to store monitors for each (class, method) pair
-    monitors_registery: Dict[Tuple[Type, str], List['MethodMonitor']] = {}
+    """
+    Monitor calls to a specific method on a class.
+
+    Multiple monitors may be attached to the same (class, method) pair.
+    All active monitors are executed AFTER the original method call.
+    """
+
+    _registry: Dict[Tuple[Type, str], List["MethodMonitor"]] = {}
 
     def __init__(
-            self, 
-            target: Type, 
-            callable: Callable[..., None],
-            monitor_args: tuple = (),
-            monitor_kwargs: dict = {},
+            self,
+            target: Type,
+            monitor_callable: Callable[..., None],
+            monitor_args: Optional[Tuple] = None,
+            monitor_kwargs: Optional[Dict[str, Any]] = None,
             *,
-            target_method: str = '__init__',
+            target_method: str | Callable = "__init__",
             active: bool = True,
-    ) -> None:
+        ) -> None:
         """
         A class to monitor method calls of a target class, triggering a handler function after the method is called.
 
@@ -24,12 +31,12 @@ def __init__(
 
         Args:
             target (Type): The target class whose method will be monitored.
-            callable (MonitorCallable): A callable to execute when the target method is called. -
+            monitor_callable (MonitorCallable): A callable to execute when the target method is called. -
                 Signature: monitor_callable(instance: object, *monitor_args, **monitor_kwargs). -
                 **warning**: sends `None` as the first arg to `MonitorCallable` if target method is `StaticMethod` !!
-            monitor_args (tuple): Positional arguments to pass to `callable` (default: empty tuple).
-            monitor_kwargs (dict): Keyword arguments to pass to `callable` (default: empty dict).
-            target_method (str): Name of the method to monitor (default: '__init__').
+            monitor_args (Optional[Tuple]): Positional arguments to pass to `callable` (default: empty tuple).
+            monitor_kwargs (Optional[Dict[str, Any]]): Keyword arguments to pass to `callable` (default: empty dict).
+            target_method (str | Callable): Name of the method to monitor or the method itself (default: '__init__').
             active (bool): Whether the monitor active initially (default: True).
 
         Example:
@@ -43,61 +50,90 @@ def __init__(
             >>> obj.my_method()  # Also calls `my_handler(obj)`
         """
         self._target = target
-        self._monitor_callable = callable
-        self._monitor_args = monitor_args
-        self._monitor_kwargs = monitor_kwargs
-        self._target_method = target_method
+        self._monitor_callable = monitor_callable
+        self._monitor_args = monitor_args or ()
+        self._monitor_kwargs = monitor_kwargs or {}
+        self._target_method = target_method if isinstance(target_method, str) else target_method.__name__
         self._active = active
 
-        # Add this Monitor to the list of Monitors for each (class, method)
-        key = self._create_registery_key()
-        if key not in self.monitors_registery:
-            self.monitors_registery[key] = []
-            self._wrap_class_method(target, self._target_method)
-
-        self.monitors_registery[key].append(self)
+        key = (self._target, self._target_method)
 
+        if key not in self._registry:
+            self._wrap_method()
+            self._registry[key] = []
 
-    def _create_registery_key(self) -> Tuple[Type, str]:
-        return (self._target, self._target_method)
-
-    def _create_original_name(self, method_name: str) -> str:
-        return f'__original_{method_name}'
+        self._registry[key].append(self)
 
     @staticmethod
-    def _is_static_method(method: staticmethod|classmethod|Callable[[Any] ,Any]) -> bool:
-        return isinstance(method, (staticmethod, classmethod))
-
-    def _wrap_class_method(self, target: Type, method_name: str) -> None:
-        """Wrap the target method to call all Monitors."""
-        original_name = self._create_original_name(method_name)
-
-        if not hasattr(target, method_name):
-            raise ValueError(f"The target class {target.__name__} does not have a method '{method_name}'.")
-
-        # Save the original method if not already saved
-        if not hasattr(target, original_name):
-            setattr(target, original_name, getattr(target, method_name))
-
-        original_method = getattr(target, original_name)
-
-        @wraps(original_method)
-        def new_method(*args, **kwargs) -> Any:
-            output = original_method(*args, **kwargs)
-
-            key = self._create_registery_key()
-            for monitor in MethodMonitor.monitors_registery.get(key, []):
-                if monitor.is_active():
-                    monitor._monitor_callable(
-                        args[0] if self._is_static_method(original_method) else None,
-                        *monitor._monitor_args, 
-                        **monitor._monitor_kwargs,
-                    )
-
-            return output
-
-        setattr(target, method_name, new_method)
-
+    def _get_method_type(target_class, method_name) -> str:
+        """Return 'instance', 'class', or 'static'."""
+        for cls in target_class.__mro__:
+            if method_name in cls.__dict__:
+                attr = cls.__dict__[method_name]
+                if isinstance(attr, staticmethod):
+                    return "static"
+                elif isinstance(attr, classmethod):
+                    return "class"
+                else:
+                    return "instance"
+        return "instance"
+
+    def _wrap_method(self) -> None:
+        if not hasattr(self._target, self._target_method):
+            raise AttributeError(
+                f"{self._target.__name__} has no method '{self._target_method}'"
+            )
+
+        # Get the original descriptor
+        original_attr = None
+        for cls in self._target.__mro__:
+            if self._target_method in cls.__dict__:
+                original_attr = cls.__dict__[self._target_method]
+                break
+
+        if original_attr is None:
+            raise AttributeError(f"Cannot find method {self._target_method}")
+
+        method_type = self._get_method_type(self._target, self._target_method)
+
+        # Extract the original function for wrapping
+        if isinstance(original_attr, (staticmethod, classmethod)):
+            original_func = original_attr.__func__
+        else:
+            original_func = original_attr
+
+        @wraps(original_func)
+        def wrapper(*args, **kwargs):
+            result = original_func(*args, **kwargs)
+
+            for monitor in self._registry.get((self._target, self._target_method), []):
+                if not monitor._active:
+                    continue
+
+                if method_type == "static":
+                    first_arg = None
+                else:
+                    first_arg = args[0]
+
+                monitor._monitor_callable(
+                    first_arg, *monitor._monitor_args, **monitor._monitor_kwargs
+                )
+
+            return result
+
+        # Reapply descriptor to preserve type
+        if method_type == "static":
+            wrapped = staticmethod(wrapper)
+        elif method_type == "class":
+            wrapped = classmethod(wrapper)
+        else:
+            wrapped = wrapper
+
+        # Store reference to original for unwrapping
+        wrapped.__methodmonitor_original__ = original_func  # type: ignore[attr-defined]
+        wrapped.__methodmonitor_wrapped__ = True  # type: ignore[attr-defined]
+
+        setattr(self._target, self._target_method, wrapped)
 
     def activate(self) -> None:
         """Activate the monitor."""
@@ -108,31 +144,39 @@ def deactivate(self) -> None:
         self._active = False
 
     def remove(self) -> None:
-        """Remove the handler and restore the original method if no monitors are left."""
-        key = self._create_registery_key()
-        if key in self.monitors_registery:
-            self.monitors_registery[key].remove(self)
-            if not self.monitors_registery[key]:
-                # Restore the original method
-                original_name = self._create_original_name(self._target_method)
-                if hasattr(self._target, original_name):
-                    setattr(self._target, self._target_method, getattr(self._target, original_name))
-                    delattr(self._target, original_name)
-
-                del self.monitors_registery[key]
-
+        """Remove the monitor and restore original method if no monitors left."""
+        key = (self._target, self._target_method)
+        monitors = self._registry.get(key)
+        if not monitors:
+            return
+
+        if self in monitors:
+            monitors.remove(self)
+
+        if not monitors:
+            wrapped = getattr(self._target, self._target_method)
+            original = getattr(wrapped, "__methodmonitor_original__", None)
+            if original:
+                # Reapply original descriptor type
+                method_type = self._get_method_type(self._target, self._target_method)
+                if method_type == "static":
+                    original = staticmethod(original)
+                elif method_type == "class":
+                    original = classmethod(original)
+                setattr(self._target, self._target_method, original)
+
+            del self._registry[key]
 
     def is_active(self) -> bool:
-        return bool(self._active)
+        return self._active
 
     def __bool__(self) -> bool:
-        return self.is_active()
-
-    def __str__(self) -> str:
-        return f'<MethodMonitor of: {self._target} (method={self._target_method})>'
+        return self._active
 
     def __repr__(self) -> str:
-        return f'MethodMonitor({self._target}, {self._monitor_callable}, target_method={self._target_method}, monitor_args={self._monitor_args}, monitor_kwargs={self._monitor_kwargs})'
-
-    def __del__(self) -> None:
-        self.remove()
+        return (
+            f"MethodMonitor("
+            f"target={self._target.__name__}, "
+            f"method={self._target_method}, "
+            f"active={self._active})"
+        )