Merge pull request #54 from ionite34/dev

docs/source/conf.py

@@ -15,7 +15,7 @@
 project = "einspect"
 copyright = "2023, Ionite"
 author = "Ionite"
-release = "v0.5.9"
+release = "v0.5.10"
 
 
 # -- General configuration ---------------------------------------------------

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "einspect"
-version = "0.5.9"
+version = "0.5.10"
 packages = [{ include = "einspect", from = "src" }]
 description = "Extended Inspect - view and modify memory structs of runtime objects."
 authors = ["ionite34 <dev@ionite.io>"]

src/einspect/__init__.py

@@ -12,6 +12,6 @@
 
 __all__ = ("view", "unsafe", "impl", "orig", "ptr", "NULL")
 
-__version__ = "0.5.9"
+__version__ = "0.5.10"
 
 unsafe = global_unsafe

src/einspect/api.py

@@ -187,6 +187,10 @@ def Realloc(p: c_void_p, n: Annotated[int, c_size_t]) -> c_void_p:
             """
 
 
+# Preload some core APIs to not rely on type hint inference.
+getattr(Py, "IncRef")
+getattr(Py, "DecRef")
+
 PyObj_FromPtr: Callable[[int], object] = _ctypes.PyObj_FromPtr
 """(Py_ssize_t ptr) -> Py_ssize_t"""
 

src/einspect/structs/py_type.py

@@ -222,6 +222,8 @@ def __init__(self, tp_new: newfunc, wrap_type: Type[_T]):
         # Cast tp_new to remove Structure binding
         self._tp_new = cast(tp_new, newfunc)
         self._type = wrap_type
+        # Store the original slot wrapper as well, for restoring
+        self._orig_slot_fn = self._type.__new__
         self.__name__ = "__new__"
 
     def __repr__(self):

src/einspect/type_orig.py

@@ -4,40 +4,46 @@
 from ctypes import cast
 from types import BuiltinFunctionType
 from typing import Any, Type, TypeVar
+from weakref import WeakKeyDictionary
 
 from einspect.structs.include.object_h import newfunc
 from einspect.structs.py_type import PyTypeObject, TypeNewWrapper
 
 _T = TypeVar("_T")
 MISSING = object()
 
+# Statically cache some methods used in cache lookups
 obj_tp_new = cast(PyTypeObject.from_object(object).tp_new, newfunc)
 obj_getattr = object.__getattribute__
 type_hash = type.__hash__
 str_eq = str.__eq__
+
 dict_setdefault = dict.setdefault
 dict_contains = dict.__contains__
-dict_get = dict.get
 dict_getitem = dict.__getitem__
+dict_get = dict.get
+
+wk_dict_setdefault = WeakKeyDictionary.setdefault
+wk_dict_getitem = WeakKeyDictionary.__getitem__
 
-_slots_cache: dict[type, dict[str, Any]] = {}
+# Cache of original type attributes, keys are weakrefs to not delay GC of user types
+_cache: WeakKeyDictionary[type, dict[str, Any]] = WeakKeyDictionary()
 
 
 def add_cache(type_: Type[_T], name: str, method: Any) -> Any:
     """Add a method to the cache."""
-    type_methods = dict_setdefault(_slots_cache, type_, {})
+    type_methods = wk_dict_setdefault(_cache, type_, {})
 
-    # For `__new__` methods, use special TypeNewWrapper to use modified safety check
+    # For `__new__` methods, use special TypeNewWrapper for modified safety check
     if name == "__new__":
         # Check if we're trying to set a previous impl method
         # If so, avoid the loop by using object.__new__
         if not isinstance(method, BuiltinFunctionType):
             method = get_cache(object, "__new__")
         else:
             tp_new = PyTypeObject.from_object(type_).tp_new
-            obj = obj_tp_new(TypeNewWrapper, (), {})
-            obj.__init__(tp_new, type_)
-            method = obj
+            method = obj_tp_new(TypeNewWrapper, (), {})
+            method.__init__(tp_new, type_)
 
     # Only allow adding once, ignore if already added
     dict_setdefault(type_methods, name, method)
@@ -46,13 +52,13 @@ def add_cache(type_: Type[_T], name: str, method: Any) -> Any:
 
 def in_cache(type_: type, name: str) -> bool:
     """Return True if the method is in the cache."""
-    type_methods = dict_setdefault(_slots_cache, type_, {})
+    type_methods = wk_dict_setdefault(_cache, type_, {})
     return dict_contains(type_methods, name)
 
 
 def get_cache(type_: type, name: str) -> Any:
     """Get the method from the type in cache."""
-    type_methods = dict_setdefault(_slots_cache, type_, {})
+    type_methods = wk_dict_setdefault(_cache, type_, {})
     return dict_getitem(type_methods, name)
 
 
@@ -71,9 +77,9 @@ class orig:
     def __new__(cls, type_: Type[_T]) -> Type[_T]:
         # To avoid a circular call loop when orig is called within
         # impl of object.__new__, we use the raw tp_new of object here.
-        obj = obj_tp_new(cls, (), {})
-        obj.__type = type_
-        return obj  # type: ignore
+        self = obj_tp_new(cls, (), {})
+        self.__type = type_
+        return self  # type: ignore
 
     def __repr__(self) -> str:
         return f"orig({self.__type.__name__})"

src/einspect/views/view_type.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import sys
+import weakref
 from collections.abc import Generator, Sequence
 from contextlib import contextmanager, suppress
 from typing import TYPE_CHECKING, Any, Callable, Literal, Type, TypeVar, Union, get_args
@@ -12,6 +13,7 @@
 from einspect.errors import UnsafeError
 from einspect.structs import PyTypeObject
 from einspect.structs.include.object_h import TpFlags
+from einspect.structs.py_type import TypeNewWrapper
 from einspect.structs.slots_map import (
     Slot,
     get_slot,
@@ -20,7 +22,7 @@
     tp_as_number,
     tp_as_sequence,
 )
-from einspect.type_orig import add_cache, in_cache
+from einspect.type_orig import add_cache, get_cache, in_cache
 from einspect.views.view_base import REF_DEFAULT, VarView
 
 if TYPE_CHECKING:
@@ -38,6 +40,20 @@
 ALLOC_MODES = frozenset({"mapping", "sequence", "all"})
 
 
+def get_func_name(func: Callable) -> str:
+    """Returns the name of the function."""
+    return get_func_base(func).__name__
+
+
+def get_func_base(func: Callable) -> Callable:
+    """Returns the base function of a method or property."""
+    if isinstance(func, property):
+        return func.fget
+    elif isinstance(func, (classmethod, staticmethod)):
+        return func.__func__
+    return func
+
+
 def _to_types(
     types_or_unions: Sequence[type | UnionType],
 ) -> Generator[type, None, None]:
@@ -51,9 +67,48 @@ def _to_types(
             raise TypeError(f"cls must be a type or Union, not {t.__class__.__name__}")
 
 
+def _restore_impl(*types: type, name: str) -> None:
+    """
+    Finalizer to restore the original `name` attribute on type(s).
+
+    If there is no original attribute, delete the attribute.
+    """
+    for t in types:
+        v = TypeView(t, ref=False)
+        # Get the original attribute from cache
+        if in_cache(t, name):
+            attr = get_cache(t, name)
+            # For TypeNewWrapper, use the original slot wrapper
+            if isinstance(t, TypeNewWrapper):
+                attr = t._orig_slot_fn
+            # Set the attribute back using a view
+            v[name] = attr
+        else:
+            # If there is no original attribute, delete the attribute
+            with v.as_mutable():
+                delattr(t, name)
+
+
+def _attach_finalizer(types: Sequence[type], func: Callable) -> None:
+    """Attaches a finalizer to the function to remove the implemented method on types."""
+    # Use the base function (we can't set attributes on properties)
+    func = get_func_base(func)
+    name = func.__name__
+    # Don't finalize types that are already registered
+    if hasattr(func, "_impl_types"):
+        types = [t for t in types if t not in func._impl_types]
+        # Update list
+        func._impl_types.extend(types)
+    else:
+        func._impl_types = list(types)
+
+    func._impl_finalize = weakref.finalize(func, _restore_impl, *types, name=name)
+
+
 def impl(
     *cls: Type[_T] | UnionType,
     alloc: AllocMode | None = None,
+    detach: bool = False,
 ) -> Callable[[_Fn], _Fn]:
     # noinspection PyShadowingNames, PyCallingNonCallable
     """
@@ -62,12 +117,19 @@ def impl(
     Supports methods decorated with property, classmethod, or staticmethod.
 
     Args:
-        cls: The types to implement the method on. Can be types or Unions.
-        alloc: The allocation type of the type. If the type is a mapping or
+        cls: The type(s) or Union(s) to implement the method on.
+        alloc: The PyMethod allocation mode. Default of None will automatically allocate
+            PyMethod structs as needed. If "sequence" or "mapping", will prefer the
+            respective PySequenceMethods or PyMappingMethods in cases of ambiguious slot names.
+            (e.g. "__getitem__" or "__len__"). If "all", will allocate all PyMethod structs.
+        detach: If True, will remove the implemented method from the type when
+            the decorated function is garbage collected. This will hold a reference to
+            the type(s) for the lifetime of the function. Requires function to support weakrefs.
 
     Returns:
         The original function after it has been implemented on the types,
         allows chaining of multiple impl decorators.
+
     Examples:
         >>> @impl(int)
         ... def is_even(self):
@@ -81,19 +143,27 @@ def impl(
         ...     except ValueError:
         ...         return None
     """
-    targets = list(_to_types(cls))
+    targets = tuple(_to_types(cls))
 
     def wrapper(func: _Fn) -> _Fn:
-        if isinstance(func, property):
-            name = func.fget.__name__
-        else:
-            name = func.__name__
+        # detach requires weakrefs
+        try:
+            weakref.ref(get_func_base(func))
+        except TypeError:
+            raise TypeError(
+                f"detach=True requires function {func!r} to support weakrefs"
+            ) from None
+
+        name = get_func_name(func)
 
         for type_ in targets:
             t_view = TypeView(type_)
             with t_view.alloc_mode(alloc):
                 t_view[name] = func
 
+        if detach:
+            _attach_finalizer(targets, func)
+
         return func
 
     return wrapper