From dacebd4a75bf74a5904d4c60eed626646966382b Mon Sep 17 00:00:00 2001
From: Thor Whalen <1906276+thorwhalen@users.noreply.github.com>
Date: Tue, 26 Nov 2024 10:35:58 +0100
Subject: [PATCH] feat: add redirect_getattr_to_getitem decorator to redirect
 attribute access to __getitem__

---
 dol/__init__.py         |  3 ++-
 dol/tests/test_trans.py | 40 ++++++++++++++++++++++++++++++++++-
 dol/trans.py            | 47 ++++++++++++++++++++++++++++++++++++++---
 3 files changed, 85 insertions(+), 5 deletions(-)

diff --git a/dol/__init__.py b/dol/__init__.py
index 294a942a..2df6dd9a 100644
--- a/dol/__init__.py
+++ b/dol/__init__.py
@@ -133,7 +133,8 @@ def ihead(store, n=1):
     insert_aliases,  # insert aliases for special (dunder) store methods,
     add_missing_key_handling,  # add a missing key handler to a store
     cache_iter,  # being deprecated
-    store_decorator,  # Helper to make store decorators.
+    store_decorator,  # Helper to make store decorators
+    redirect_getattr_to_getitem,  # redirect attribute access to __getitem__ 
 )
 
 from dol.caching import (
diff --git a/dol/tests/test_trans.py b/dol/tests/test_trans.py
index 7635b750..f0e98174 100644
--- a/dol/tests/test_trans.py
+++ b/dol/tests/test_trans.py
@@ -1,6 +1,6 @@
 """Test trans.py functionality."""
 
-from dol.trans import filt_iter
+from dol.trans import filt_iter, redirect_getattr_to_getitem
 
 
 def test_filt_iter():
@@ -24,3 +24,41 @@ def test_filt_iter():
     d = {"test.txt": 1, "report.doc": 2, "image.jpg": 3}
     dd = is_text(d)
     assert dict(dd) == {"test.txt": 1, "report.doc": 2}
+
+
+def test_redirect_getattr_to_getitem():
+
+    # Applying it to a class
+
+    ## ... with the @decorator syntax
+    @redirect_getattr_to_getitem
+    class MyDict(dict):
+        pass
+
+    d1 = MyDict(a=1, b=2)
+    assert d1.a == 1
+    assert d1.b == 2
+    assert list(d1) == ['a', 'b']
+
+    ## ... as a decorator factory
+    D = redirect_getattr_to_getitem()(dict)
+    d2 = D(a=1, b=2)
+    assert d2.a == 1
+    assert d2.b == 2
+    assert list(d2) == ['a', 'b']
+
+    # Applying it to an instance
+
+    ## ... as a decorator
+    backend_d = dict(a=1, b=2)
+
+    d3 = redirect_getattr_to_getitem(backend_d)
+    assert d3.a == 1
+    assert d3.b == 2
+    assert list(d3) == ['a', 'b']
+
+    ## ... as a decorator factory
+    d4 = redirect_getattr_to_getitem()(backend_d)
+    assert d4.a == 1
+    assert d4.b == 2
+    assert list(d4) == ['a', 'b']
diff --git a/dol/trans.py b/dol/trans.py
index 594c015f..14b86255 100644
--- a/dol/trans.py
+++ b/dol/trans.py
@@ -347,9 +347,9 @@ def _func_wrapping_store_in_cls_if_not_type(store, **kwargs):
 
         return r
 
-    _func_wrapping_store_in_cls_if_not_type.func = (
-        func  # TODO: look for usages, and if not, use __wrapped__
-    )
+    # Two standard attributes for storing the original function are func and __wrapped__
+    _func_wrapping_store_in_cls_if_not_type.func = func
+    _func_wrapping_store_in_cls_if_not_type.__wrapped__ = func
 
     # @wraps(func)
     wrapper_sig = Sig(func).merge_with_sig(
@@ -3344,3 +3344,44 @@ def affix_key_codec(prefix: str = "", suffix: str = ""):
         encoder=partial(_affix_encoder, prefix=prefix, suffix=suffix),
         decoder=partial(_affix_decoder, prefix=prefix, suffix=suffix),
     )
+
+
+@store_decorator
+def redirect_getattr_to_getitem(cls=None, *, keys_have_priority_over_attributes=False):
+    """A mapping decorator that redirects attribute access to __getitem__.
+
+    Warning: This decorator will make your class un-pickleable.
+
+    :param keys_have_priority_over_attributes: If True, keys will have priority over existing attributes.
+
+    >>> @redirect_getattr_to_getitem
+    ... class MyDict(dict):
+    ...     pass
+    >>> d = MyDict(a=1, b=2)
+    >>> d.a
+    1
+    >>> d.b
+    2
+    >>> list(d)
+    ['a', 'b']
+
+    """
+
+    class RidirectGetattrToGetitem(cls):
+        """A class that redirects attribute access to __getitem__"""
+
+        _keys_have_priority_over_attributes = keys_have_priority_over_attributes
+
+        def __getattr__(self, attr):
+            if attr in self:
+                if self._keys_have_priority_over_attributes or attr not in dir(
+                    type(self)
+                ):
+                    return self[attr]
+            # if attr not in self, or if it is in the class, then do normal getattr
+            return super(RidirectGetattrToGetitem, self).__getattr__(attr)
+
+        def __dir__(self) -> Iterable[str]:
+            return list(self)
+
+    return RidirectGetattrToGetitem