From 0c0ceecf17de1f8ae70838db869b6d69f4f2ccc1 Mon Sep 17 00:00:00 2001
From: Matthew Wardrop <mpwardrop@gmail.com>
Date: Thu, 15 Aug 2024 15:29:10 -0700
Subject: [PATCH] Extend `Alias` with support for nested structure.

---
 spec_classes/types/alias.py    | 93 ++++++++++++++++++++++++++++------
 tests/types/test_alias.py      | 49 ++++++++++++++++--
 tests/types/test_attr_proxy.py | 10 +++-
 3 files changed, 131 insertions(+), 21 deletions(-)

diff --git a/spec_classes/types/alias.py b/spec_classes/types/alias.py
index 0f0da87..a1b4463 100644
--- a/spec_classes/types/alias.py
+++ b/spec_classes/types/alias.py
@@ -1,6 +1,11 @@
+import ast
+import functools
+import re
 import warnings
 from typing import Any, Callable, Optional, Type
 
+from cached_property import cached_property
+
 from .missing import MISSING
 
 
@@ -27,7 +32,11 @@ class Alias:
     retrieve it.
 
     Attributes:
-        attr: The name of the attribute to be aliased.
+        attr: The name of the attribute to be aliased. If `attr` has periods in
+            it (e.g. `foo.bar`), then the attribute will be looked up
+            transitively (e.g. `self.foo.bar`). This also works through
+            mappings, so if `attr` is `foo["bar"]`, than bar will be looked up
+            using `self.foo["bar"]`.
         passthrough: Whether to pass through mutations of the `Alias`
             attribute through to the aliased attribute. (default: False)
         transform: An optional unary transform to apply to the value of the
@@ -38,6 +47,10 @@ class Alias:
             the underlying aliased attribute value are passed through).
     """
 
+    ATTR_PARSER = re.compile(
+        r"(?:(?<!^)\.)?(?P<lookup>(?<!\.)\[\"(?:[^\"\\]|\\.)*\"\](?!\.)|(?<!\.)\['(?:[^'\\](?!\.)|\\.)*'\]|\w+)"
+    )
+
     def __init__(
         self,
         attr: str,
@@ -54,23 +67,62 @@ def __init__(
         self._owner = None
         self._owner_attr = None
 
+    def __setattr__(self, attr, value):
+        # Ensure that the _attr_path is kept up to date.
+        super().__setattr__(attr, value)
+        if attr == "attr":
+            self.__dict__.pop("_attr_path", None)
+            self._attr_path
+
+    @cached_property
+    def _attr_path(self):
+        if self.attr.isidentifier():
+            return [self.attr]
+        matches = list(self.ATTR_PARSER.finditer(self.attr))
+        if not "".join(match.group(0) for match in matches) == self.attr:
+            raise ValueError(f"Invalid attribute path: {self.attr}")
+        return [match.group("lookup") for match in matches]
+
     @property
     def override_attr(self):
-        if not self._owner_attr and not self.passthrough:
+        if self.passthrough:
             return None
-        return (
-            self.attr
-            if self.passthrough
-            else f"__spec_classes_Alias_{self._owner_attr}_override"
-        )
+        if not self._owner_attr:
+            raise RuntimeError(
+                f"`{self.__class__.__name__}` instances must be assigned to a "
+                "class attribute before they can be used."
+            )
+        return f"__spec_classes_Alias_{self._owner_attr}_override"
+
+    def __lookup_attr_path(self, instance, attr_path):
+        try:
+            return functools.reduce(
+                lambda obj, attr: (
+                    obj[ast.literal_eval(attr[1:-1])]
+                    if attr.startswith("[")
+                    else getattr(obj, attr)
+                ),
+                attr_path,
+                instance,
+            )
+        except (AttributeError, KeyError) as e:
+            raise AttributeError(
+                f"`{instance.__class__.__name__}{'' if self.attr.startswith('[') else '.'}{self.attr}` [Caused by: {e}]"
+            ) from e
 
     def __get__(self, instance: Any, owner=None):
         if instance is None:
             return self
-        if self._owner_attr and hasattr(instance, self.override_attr):
+        if (
+            self._owner_attr
+            and not self.passthrough
+            and hasattr(instance, self.override_attr)
+        ):
             return getattr(instance, self.override_attr)
         try:
-            return (self.transform or (lambda x: x))(getattr(instance, self.attr))
+            return (self.transform or (lambda x: x))(
+                self.__lookup_attr_path(instance, self._attr_path)
+            )
         except AttributeError:
             if self.fallback is not MISSING:
                 return self.fallback
@@ -83,15 +135,24 @@ def __get__(self, instance: Any, owner=None):
             ) from e
 
     def __set__(self, instance, value):
-        if not self.override_attr:
-            raise RuntimeError(
-                f"Attempting to set the value of an `{self.__class__.__name__}` "
-                "instance that is not properly associated with a class."
-            )
-        setattr(instance, self.override_attr, value)
+        if self.passthrough:
+            obj = self.__lookup_attr_path(instance, self._attr_path[:-1])
+            if self._attr_path[-1].startswith("["):
+                obj[ast.literal_eval(self._attr_path[-1][1:-1])] = value
+            else:
+                setattr(obj, self._attr_path[-1], value)
+        else:
+            setattr(instance, self.override_attr, value)
 
     def __delete__(self, instance):
-        delattr(instance, self.override_attr)
+        if self.passthrough:
+            obj = self.__lookup_attr_path(instance, self._attr_path[:-1])
+            if self._attr_path[-1].startswith("["):
+                del obj[ast.literal_eval(self._attr_path[-1][1:-1])]
+            else:
+                delattr(obj, self._attr_path[-1])
+        else:
+            delattr(instance, self.override_attr)
 
     def __set_name__(self, owner, name):
         self._owner = owner
diff --git a/tests/types/test_alias.py b/tests/types/test_alias.py
index 46e9a48..9ee6fba 100644
--- a/tests/types/test_alias.py
+++ b/tests/types/test_alias.py
@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import re
 
 import pytest
 
-from spec_classes import MISSING, Alias, DeprecatedAlias, spec_class
+from spec_classes import MISSING, Alias, DeprecatedAlias, spec_class, spec_property
 
 
 def test_alias():
@@ -15,6 +17,17 @@ class Item:
         v: int = Alias("x", passthrough=True, fallback=2)
         u: int = Alias("u")
 
+        subitem: Item
+        data: dict = {"Hello": "World"}
+
+        @spec_property(cache=True)
+        def subitem(self) -> Item:
+            return Item(x=10)
+
+        r: str = Alias("data['Hello']", passthrough=True)
+        s: int = Alias("subitem.x", passthrough=True)
+        t: int = Alias("subitem.x")
+
     assert Item(x=2).y == 2
     assert Item(x=2).z == 4
     assert Item().w == 2
@@ -58,16 +71,46 @@ class Item:
     ):
         Item().y
 
-    assert Alias("attr").override_attr is None
+    with pytest.raises(
+        RuntimeError,
+        match=re.escape(
+            "`Alias` instances must be assigned to a class attribute before they can be used."
+        ),
+    ):
+        Alias("attr").override_attr
 
     with pytest.raises(
         RuntimeError,
         match=re.escape(
-            "Attempting to set the value of an `Alias` instance that is not properly associated with a class."
+            "`Alias` instances must be assigned to a class attribute before they can be used."
         ),
     ):
         Alias("attr").__set__(None, "Hi")
 
+    for path in ["a.", ".b", "c[]", "c[[", "c.['d']"]:
+        with pytest.raises(
+            ValueError, match=re.escape(f"Invalid attribute path: {path}")
+        ):
+            Alias(path)
+
+    # Test Subitems
+    assert Item().s == 10
+    assert Item().t == 10
+    item = Item()
+    item.t = 20
+    assert item.subitem.x == 10
+    assert item.t == 20
+    assert item.s == 10
+    item.s = 20
+    assert item.subitem.x == 20
+    assert item.r == "World"
+    item.r = "World!"
+    assert item.data["Hello"] == "World!"
+    del item.r
+    assert "Hello" not in item.data
+
+    assert Item.r.override_attr is None
+
 
 def test_deprecated_alias():
     @spec_class
diff --git a/tests/types/test_attr_proxy.py b/tests/types/test_attr_proxy.py
index 081ba3c..926dc37 100644
--- a/tests/types/test_attr_proxy.py
+++ b/tests/types/test_attr_proxy.py
@@ -59,12 +59,18 @@ class Item:
     ):
         Item().y
 
-    assert AttrProxy("attr").override_attr is None
+    with pytest.raises(
+        RuntimeError,
+        match=re.escape(
+            "`AttrProxy` instances must be assigned to a class attribute before they can be used."
+        ),
+    ):
+        AttrProxy("attr").override_attr
 
     with pytest.raises(
         RuntimeError,
         match=re.escape(
-            "Attempting to set the value of an `AttrProxy` instance that is not properly associated with a class."
+            "`AttrProxy` instances must be assigned to a class attribute before they can be used."
         ),
     ):
         AttrProxy("attr").__set__(None, "Hi")