Merge pull request #85 from ionite34/fix-typeview-restore

Author: ionite34

docs/source/conf.py

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

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "einspect"
-version = "0.5.15"
+version = "0.5.16"
 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

@@ -8,7 +8,7 @@
 from einspect.views.view_type import impl
 
 __all__ = ("view", "unsafe", "impl", "orig", "ptr", "NULL", "__version__")
-__version__ = "0.5.15"
+__version__ = "0.5.16"
 
 unsafe = global_unsafe
 

src/einspect/type_orig.py

@@ -66,6 +66,14 @@ def add_impls(type_: type, *attrs: str) -> None:
     attrs_set.update(attrs)
 
 
+def remove_impls(type_: type, *attrs: str) -> None:
+    """Remove a set of implemented attributes from the cache."""
+    attrs_set = wk_dict_getitem(_impls, type_)
+    if attrs_set is not None:
+        for attr in attrs:
+            attrs_set.discard(attr)
+
+
 def try_cache_attr(
     type_: type,
     name: str,
@@ -133,6 +141,16 @@ def get_type_cache(type_: type) -> dict[str, Any]:
         ) from None
 
 
+def get_impls(type_: type) -> set[str]:
+    """Get the impls cache for the type."""
+    try:
+        return wk_dict_getitem(_impls, type_)
+    except KeyError:
+        raise KeyError(
+            f"Original attributes cache was not found for type {type_!r}"
+        ) from None
+
+
 def normalize_slot_attr(attr: object | TypeNewWrapper) -> Any:
     """Normalize a slot attribute to its original value."""
     if isinstance(attr, TypeNewWrapper):

src/einspect/views/view_type.py

@@ -24,10 +24,11 @@
     MISSING,
     add_impls,
     get_cache,
-    get_type_cache,
+    get_impls,
     in_cache,
     in_impls,
     normalize_slot_attr,
+    remove_impls,
     try_cache_attr,
 )
 from einspect.views.view_base import REF_DEFAULT, VarView
@@ -352,26 +353,25 @@ def restore(self, *names: str | Callable) -> None:
 
         if not names:
             # Restore all attributes
-            type_cache = get_type_cache(type_)
-            for name, attr in type_cache.items():
-                self[name] = normalize_slot_attr(attr)
-            return type_
+            names = get_impls(type_)
 
         # Normalize names of stuff like properties and class methods
         names = [get_func_name(n) if callable(n) else n for n in names]
         # Check all names exist first
         for name in names:
-            if in_cache(type_, name):
+            if in_impls(type_, name) and in_cache(type_, name):
                 attr = get_cache(type_, name)
                 # MISSING is a special case, it means there is no original attribute
                 if attr is MISSING:
                     # Same as no attribute, delete it
                     del self[name]
+                    remove_impls(type_, name)
                 else:
                     self[name] = normalize_slot_attr(attr)
             # If in impl record, and not in cache, remove the attribute
             elif in_impls(type_, name):
                 del self[name]
+                remove_impls(type_, name)
             else:
                 raise AttributeError(
                     f"{type_.__name__!r} has no original attribute {name!r}"

tests/test_impl_sp.py

@@ -4,6 +4,38 @@
 import pytest
 
 from einspect import impl, orig, view
+from einspect.type_orig import _impls
+
+
+@pytest.mark.run_in_subprocess
+def test_impl_restore():
+    # Clear impls cache
+    backup_impls = _impls.copy()
+    _impls.clear()
+
+    @impl(int)
+    def __repr__(self):
+        return "repr"
+
+    @impl(int)
+    def x(self):
+        return "x"
+
+    n = 5
+    assert repr(n) == "repr"
+    # noinspection PyUnresolvedReferences
+    assert n.x() == "x"
+
+    view(int).restore()
+
+    assert repr(n) == "5"
+    with pytest.raises(AttributeError):
+        # noinspection PyUnresolvedReferences
+        _ = n.x()
+
+    # Restore impls cache
+    _impls.clear()
+    _impls.update(backup_impls)
 
 
 @pytest.mark.run_in_subprocess