Add wrapped_self() — Phase 1 fix for #18 (self is unwrapped in delegation-wrapped stores)

CLAUDE.md

@@ -215,7 +215,7 @@ Run tests: `pytest dol/tests/`
 
 ## Known Limitations / Gotchas
 
-- **`wrap_kvs` + `self` inside methods**: When a `wrap_kvs`-decorated class uses `self[k]` in its own methods, `self` is the unwrapped instance. Re-apply the wrapper to `self` if transforms are needed (Issue #18, still open — delegation architecture).
+- **`wrap_kvs` + `self` inside methods**: When a `wrap_kvs`-decorated class uses `self[k]` in its own methods, `self` is the unwrapped inner store, so transforms are bypassed (Issue #18 — delegation architecture). Blessed fix: `from dol import wrapped_self` and write `wrapped_self(self)[k]` to reach the outer, transform-applying store (climbs to the outermost wrapper for stacked/`Pipe` wraps; a safe no-op on direct `Store` subclasses). The older `sq(self)[k]` re-wrap still works. A structural fix (is-a wrapping) is proposed in `misc/docs/dol_issue18_design.md`.
 - **`clear()` is disabled** on `KvPersister`. Call `ensure_clear_to_kv_store(store)` to re-enable.
 - **No async support** in core. Use synchronous wrappers for async backends (thread pool, etc.).
 - **Transforms wanting the store**: a transform is called `f(self, data)` only if its first param is named `self`/`store`/`mapping` **and** it has ≥2 required params; otherwise `f(data)`. Mark explicitly with `FirstArgIsMapping(f)`. (`bytes.decode` as `obj_of_data` now works — Issue #9 fixed.)

dol/__init__.py

@@ -57,6 +57,7 @@ def ihead(store, n=1):
     Persister,  # TODO: deprecate? (now KvPersister)
     kv_walk,  # walk a kv store
     Store,  # base class for stores (adds hooks for key and value transforms)
+    wrapped_self,  # recover the outer (transform-applying) store inside a wrapped method
 )
 
 

dol/base.py

@@ -26,6 +26,7 @@
 
 from functools import partial, update_wrapper
 import copyreg
+import weakref
 from collections.abc import Collection as CollectionABC
 from collections.abc import Mapping, MutableMapping
 from collections.abc import (
@@ -262,17 +263,17 @@ def __get__(self, instance, owner):
             #   self.__wrapped__ would make it hard to debug and
             #   self would fail with unbound methods (why?)
             #   So doing a check here, but would like to find a better solution.
-            wrapped_self = getattr(self, "__wrapped__", None)
-            if is_classmethod(wrapped_self) or is_unbound_method(wrapped_self):
-                return wrapped_self
+            wrapped_callable = getattr(self, "__wrapped__", None)
+            if is_classmethod(wrapped_callable) or is_unbound_method(wrapped_callable):
+                return wrapped_callable
             else:
                 return self
 
-            # wrapped_self = getattr(self, '__wrapped__', None)
-            # if not is_classmethod(wrapped_self):
+            # wrapped_callable = getattr(self, '__wrapped__', None)
+            # if not is_classmethod(wrapped_callable):
             #     return self
             # else:
-            #     return wrapped_self
+            #     return wrapped_callable
         else:
             # i.e. return instance.delegate.attr
             return getattr(getattr(instance, self.delegate_name), self.attr_name)
@@ -285,6 +286,130 @@ def __delete__(self, instance):
         delattr(getattr(instance, self.delegate_name), self.attr_name)
 
 
+# ---------------------------------------------------------------------------------------
+# wrapped_self: recover the outer (transform-applying) store from inside a wrapped
+# class's own method. See Issue #18 and misc/docs/dol_issue18_design.md.
+#
+# dol wraps by delegation (has-a): ``wrap_kvs(SomeClass)`` returns a ``Store`` subclass
+# ``Wrap`` that HOLDS an instance of ``SomeClass`` in ``self.store``. A method defined on
+# ``SomeClass`` is copied onto ``Wrap`` as a ``DelegatedAttribute`` and runs bound to the
+# INNER, unwrapped ``self.store`` -- so ``self[k]`` inside such a method bypasses the
+# transform pipeline (Issue #18). ``wrapped_self(self)`` maps that inner store back to the
+# outer wrapper so ``self[k]`` can be routed through the transforms as
+# ``wrapped_self(self)[k]``.
+#
+# The mapping is kept in a module-global registry of weak references, keyed by ``id`` of
+# the inner store, populated in ``Store.__init__`` and re-populated in ``Store.__setstate__``
+# (unpickling bypasses ``__init__``). We key by ``id`` -- rather than setting an attribute on
+# the inner store -- because inner stores are often plain ``dict``/C backends that reject
+# instance attributes and whose pickling a stray weakref would break.
+#
+# Each id maps to a *list* of weakrefs, not a single one, so that objects legitimately
+# sharing an inner store (most importantly a shallow ``copy.copy`` of a wrapped store, whose
+# copy shares ``self.store`` with the original) coexist: a finalizer removes only its own
+# weakref, never a live sibling's. ``wrapped_self`` then picks a live wrapper that still
+# actually holds the queried object as its ``.store`` (an identity guard).
+_wrapper_backrefs: "dict[int, list[weakref.ReferenceType]]" = {}
+
+
+def _register_wrapper_backref(inner: Any, wrapper: "Store") -> None:
+    """Record that ``wrapper`` holds ``inner`` as its delegated store.
+
+    Appends ``weakref(wrapper)`` to the list registered under ``id(inner)`` (idempotent for
+    an already-registered live wrapper). When ``wrapper`` is garbage-collected the weakref
+    callback removes *only its own* entry from that list, and drops the ``id(inner)`` key once
+    the list is empty -- so garbage-collecting one wrapper (e.g. a transient ``copy.copy``)
+    never evicts a still-live sibling that shares the same inner.
+    """
+    inner_id = id(inner)
+
+    existing = _wrapper_backrefs.get(inner_id)
+    if existing is not None and any(ref() is wrapper for ref in existing):
+        return  # this live wrapper is already registered for this inner
+
+    def _cleanup(dead_ref: "weakref.ReferenceType") -> None:
+        refs = _wrapper_backrefs.get(inner_id)
+        if refs is None:
+            return
+        try:
+            refs.remove(dead_ref)
+        except ValueError:
+            pass
+        if not refs:
+            _wrapper_backrefs.pop(inner_id, None)
+
+    try:
+        ref = weakref.ref(wrapper, _cleanup)
+    except TypeError:
+        # wrapper not weak-referenceable (not expected for a Store) -- skip silently:
+        # wrapped_self simply falls back to returning the inner store unchanged.
+        return
+    _wrapper_backrefs.setdefault(inner_id, []).append(ref)
+
+
+def wrapped_self(obj: Any) -> Any:
+    """Return the outermost transform-applying store wrapping ``obj``, else ``obj`` itself.
+
+    Use this inside a method DEFINED ON a class that was wrapped by dol's delegation
+    machinery (``wrap_kvs`` / ``filt_iter`` / ``cached_keys`` / ``mk_relative_path_store`` /
+    ``Store.wrap`` applied to the *class*). There, ``self`` is bound to the inner, unwrapped
+    store, so ``self[k]`` bypasses the transforms (Issue #18). Writing
+    ``wrapped_self(self)[k]`` recovers the outer, transform-applying view.
+
+    On a direct ``Store``/``KvReader`` subclass, or a plain object that never went through
+    the delegation machinery, there is no registered wrapper and ``obj`` is returned
+    unchanged -- a safe no-op. If the store was wrapped in several layers (e.g. via
+    ``Pipe``), this climbs to the OUTERMOST wrapper.
+
+    Limitation: if the *same inner store instance* is wrapped by several live wrappers with
+    different transforms (e.g. calling ``wrap_kvs(shared_instance, ...)`` twice), a method
+    bound to that shared inner cannot know which wrapper it was reached through, so one of
+    the wrappers is returned (ambiguous but never raw/None). This does not arise for
+    class-wraps (each construction builds its own inner) nor for ``copy.copy`` (whose copies
+    are transform-equivalent, so either answer is correct).
+
+    >>> import math
+    >>> from dol import wrap_kvs, wrapped_self
+    >>> sq = wrap_kvs(data_of_obj=lambda x: x * x, obj_of_data=lambda x: math.sqrt(x))
+    >>> @sq
+    ... class S(dict):
+    ...     def via_self(self, k):
+    ...         return self[k]                 # self is the INNER store: NOT transformed
+    ...     def via_wrapped_self(self, k):
+    ...         return wrapped_self(self)[k]    # outer store: transform applied
+    >>> s = S()
+    >>> s['2'] = 2
+    >>> s['2']                                 # external access is transformed
+    2.0
+    >>> s.via_self('2')                        # Issue #18: bypasses the transform
+    4
+    >>> s.via_wrapped_self('2')                # wrapped_self recovers the transformed value
+    2.0
+    """
+    seen = {id(obj)}
+    cur = obj
+    while True:
+        nxt = None
+        for ref in _wrapper_backrefs.get(id(cur), ()):
+            candidate = ref()
+            # Identity guard: only follow a back-reference whose wrapper STILL holds ``cur``
+            # as its ``.store``. This rejects stale entries left by a post-construction
+            # ``.store`` reassignment (and the id-reuse false positives they enable), and
+            # picks a genuinely-wrapping candidate when several are registered for one id.
+            if (
+                candidate is not None
+                and id(candidate) not in seen
+                and getattr(candidate, "store", None) is cur
+            ):
+                nxt = candidate
+                break
+        if nxt is None:
+            break
+        seen.add(id(nxt))
+        cur = nxt
+    return cur
+
+
 Decorator = Callable[[Callable], Any]  # TODO: Look up typing protocols
 
 
@@ -588,6 +713,10 @@ def __init__(self, store=dict):
 
         self.store = store
 
+        # Record inner-store -> self, so wrapped_self(inner) can recover this wrapper from
+        # inside a delegated method whose ``self`` is the inner store (Issue #18).
+        _register_wrapper_backref(store, self)
+
         if hasattr(self.store, "KeysView"):
             self.KeysView = self.store.KeysView
 
@@ -731,6 +860,12 @@ def __setstate__(self, state: dict):
         for attr in Store._state_attrs:
             if attr in state:
                 setattr(self, attr, state[attr])
+        # Unpickling reconstructs via copyreg (bypassing __init__), so re-register the
+        # inner-store -> self back-reference here too (else wrapped_self would silently
+        # revert to the Issue #18 raw behavior after a pickle round-trip).
+        inner = getattr(self, "store", None)
+        if inner is not None:
+            _register_wrapper_backref(inner, self)
 
 
 # Store.register(dict)  # TODO: Would this be a good idea? To make isinstance({}, Store) be True (though missing head())

dol/tests/base_test.py

@@ -1,5 +1,7 @@
 """Testing base.py objects"""
 
+import math
+import pickle
 from typing import KT, VT, Tuple
 from collections.abc import Iterable
 import pytest
@@ -9,6 +11,7 @@
     wrap_kvs,
     filt_iter,
     cached_keys,
+    wrapped_self,
 )
 from dol.base import BaseItemsView, BaseKeysView, BaseValuesView
 from dol.trans import take_everything
@@ -195,3 +198,177 @@ def hi():
     if errors:
         first_error, *_ = errors
         raise first_error
+
+
+# ---------------------------------------------------------------------------------------
+# wrapped_self / Issue #18 ("self is the unwrapped inner store" in delegation-wrapped
+# classes). See misc/docs/dol_issue18_design.md.
+#
+# Module-level named transforms + a non-rebound base class so the wrapped store is
+# picklable (lambdas and name-rebound classes are not — pre-existing dol limitations).
+
+def _square(x):
+    return x * x
+
+
+def _sqrt(x):
+    return math.sqrt(x)
+
+
+_square_wrap = wrap_kvs(data_of_obj=_square, obj_of_data=_sqrt)
+
+
+class _SquareStoreBase(dict):
+    """A store whose stored values are the square of the value written."""
+
+    def via_self(self, k):
+        return self[k]  # self is the INNER store -> Issue #18: NOT transformed
+
+    def via_wrapped_self(self, k):
+        return wrapped_self(self)[k]  # outer store -> transform applied
+
+
+# NOTE: assign to a NEW name (do not rebind _SquareStoreBase) so the __reduce__ pickle
+# path can still resolve the original wrapped class by qualified name.
+_SquareStore = _square_wrap(_SquareStoreBase)
+
+
+def test_issue18_reproduced_and_wrapped_self_fixes_it():
+    """The #18 bug is real, and wrapped_self(self)[k] recovers the transformed value."""
+    s = _SquareStore()
+    s["2"] = 2
+    assert s["2"] == 2.0  # external access is transformed (sqrt of stored 4)
+    assert s.via_self("2") == 4  # Issue #18: inner-bound self bypasses the transform
+    assert s.via_wrapped_self("2") == 2.0  # wrapped_self recovers the transformed value
+
+
+def test_wrapped_self_climbs_to_outermost_when_nested():
+    """Under stacked/Pipe wrapping, wrapped_self returns the OUTERMOST wrapper."""
+    outer_wrap = wrap_kvs(obj_of_data=lambda x: x + 100, data_of_obj=lambda x: x - 100)
+
+    @outer_wrap
+    @_square_wrap
+    class Nested(dict):
+        def via_wrapped_self(self, k):
+            return wrapped_self(self)[k]
+
+    s = Nested()
+    s["2"] = 2
+    # via_wrapped_self must match the fully-transformed external read, not a partial one
+    assert s.via_wrapped_self("2") == s["2"]
+
+
+def test_wrapped_self_is_noop_on_direct_store_subclass():
+    """A direct Store subclass already has self == the store, so wrapped_self is a no-op."""
+
+    class Direct(Store):
+        def _obj_of_data(self, data):
+            return data * 10
+
+        def via_wrapped_self(self, k):
+            return wrapped_self(self)[k]
+
+    d = Direct()
+    d.store["a"] = 5
+    assert d["a"] == 50
+    assert d.via_wrapped_self("a") == 50  # wrapped_self(self) is self here
+    assert wrapped_self(d) is d
+
+
+def test_wrapped_self_is_noop_on_plain_objects():
+    class Plain:
+        pass
+
+    p = Plain()
+    assert wrapped_self(p) is p
+    assert wrapped_self(123) == 123
+    assert wrapped_self("hello") == "hello"
+
+
+def test_wrapped_self_covers_instance_wrap():
+    """wrap_kvs(instance) also registers (via Store.__init__), so wrapped_self resolves."""
+    inner = {"a": 1}
+    s = wrap_kvs(inner, obj_of_data=lambda x: x * 10)
+    assert s["a"] == 10
+    assert wrapped_self(s.store) is s  # inner store climbs back to the wrapper
+
+
+def test_wrapped_self_survives_pickle_round_trip():
+    """__setstate__ re-registers the backref (unpickling bypasses __init__)."""
+    s = _SquareStore()
+    s["2"] = 2
+    assert s.via_wrapped_self("2") == 2.0
+
+    s2 = pickle.loads(pickle.dumps(s))
+    assert s2["2"] == 2.0
+    # Without the __setstate__ re-registration this would silently revert to 4 (Issue #18):
+    assert s2.via_wrapped_self("2") == 2.0
+
+
+def test_wrapped_self_registration_is_zero_behavior_change():
+    """A store that never calls wrapped_self behaves exactly as before."""
+    s = wrap_kvs(dict(), key_of_id=str.upper, id_of_key=str.lower)
+    s["foo"] = 1
+    s["bar"] = 2
+    # key_of_id=str.upper => iterating yields upper-cased keys; id_of_key=str.lower on write
+    assert dict(s.items()) == {"FOO": 1, "BAR": 2}
+    assert list(s) == ["FOO", "BAR"]
+    assert "foo" in s and "nope" not in s  # __contains__ lowercases the key on lookup
+    assert len(s) == 2
+
+
+def test_wrapped_self_copy_copy_preserves_original():
+    """A shallow copy shares the inner store; GC of the copy must NOT break the original.
+
+    Regression for the single-slot id-registry clobber found in adversarial review: the
+    copy re-registered id(inner) then its finalizer popped the entry, silently reverting the
+    still-alive original to raw Issue #18 behavior.
+    """
+    import copy
+    import gc
+
+    s = _SquareStore()
+    s["2"] = 2
+    c = copy.copy(s)
+    assert c.store is s.store  # shallow copy shares the inner store
+    assert c.via_wrapped_self("2") == 2.0
+    assert s.via_wrapped_self("2") == 2.0
+    del c
+    gc.collect()
+    assert s.via_wrapped_self("2") == 2.0  # would have been 4 before the multi-valued fix
+
+
+def test_wrapped_self_deepcopy_is_independent():
+    import copy
+
+    s = _SquareStore()
+    s["2"] = 2
+    d = copy.deepcopy(s)
+    assert d.store is not s.store  # deepcopy builds a fresh inner
+    assert d.via_wrapped_self("2") == 2.0
+    assert s.via_wrapped_self("2") == 2.0
+
+
+def test_wrapped_self_registry_does_not_leak():
+    import gc
+    from dol.base import _wrapper_backrefs
+
+    gc.collect()
+    before = len(_wrapper_backrefs)
+    for _ in range(1000):
+        t = _SquareStore()
+        t["1"] = 1
+        assert t.via_wrapped_self("1") == 1.0
+    gc.collect()
+    after = len(_wrapper_backrefs)
+    assert after - before <= 5  # short-lived wrappers' entries are cleaned up
+
+
+def test_wrapped_self_shared_inner_returns_a_valid_wrapper():
+    """Two live wrappers over one shared inner: ambiguous, but never returns the raw inner."""
+    base = {}
+    a = wrap_kvs(base, obj_of_data=lambda x: x * 2, data_of_obj=lambda x: x // 2)
+    b = wrap_kvs(base, obj_of_data=lambda x: x * 3, data_of_obj=lambda x: x // 3)
+    resolved = wrapped_self(base)
+    assert resolved is a or resolved is b  # a real wrapper...
+    assert resolved is not base  # ...never the raw inner