Repository: mit-ll-responsible-ai/hydra-zen
Pull request #569 (2023-11-07 17:23:00)

Add features to `ZenStore`
Adds the following methods to `ZenStore`
- `copy`
- `copy_with_mapped_groups`
- `delete_entry` (or, `__delitem__`)
- `enqueue_all`
- `update` (or, `__ior__`)
- `merge` (or, `__or__`)
- `__len__`


Patch:
diff --git a/docs/source/changes.rst b/docs/source/changes.rst
index 5a4e3142..6252fc26 100644
--- a/docs/source/changes.rst
+++ b/docs/source/changes.rst
@@ -68,6 +68,7 @@ Improvements
 - :func:`~hydra_zen.builds` and :func:`~hydra_zen.make_custom_builds_fn` now accept a `zen_exclude` field for excluding parameters from auto-population, either by name or by pattern. See :pull:`558`.
 - :func:`~hydra_zen.builds` and :func:`~hydra_zen.just` can now configure static methods. Previously the incorrect ``_target_`` would be resolved. See :pull:`566`
 - Adds formal support for Python 3.12. See :pull:`555`
+- Several new methods were added to :class:`~hydra_zen.ZenStore`, including the abilities to copy, update, and merge stores. As well as remap the groups of a store's entries and delete individual entries. See :pull:`569`
 
 
 .. _v0.11.0:
diff --git a/docs/source/generated/hydra_zen.ZenStore.rst b/docs/source/generated/hydra_zen.ZenStore.rst
index e89375df..fbf110d6 100644
--- a/docs/source/generated/hydra_zen.ZenStore.rst
+++ b/docs/source/generated/hydra_zen.ZenStore.rst
@@ -12,8 +12,14 @@ hydra\_zen.ZenStore
    .. automethod:: __call__
    .. automethod:: __getitem__
    .. automethod:: add_to_hydra_store
+   .. automethod:: copy
+   .. automethod:: copy_with_mapped_groups
    .. automethod:: get_entry
+   .. automethod:: delete_entry
+   .. automethod:: update
+   .. automethod:: merge
    .. automethod:: has_enqueued
+   .. automethod:: enqueue_all
    .. automethod:: __iter__
    .. automethod:: __eq__
    
diff --git a/src/hydra_zen/wrapper/_implementations.py b/src/hydra_zen/wrapper/_implementations.py
index 6b20e7b7..12616e22 100644
--- a/src/hydra_zen/wrapper/_implementations.py
+++ b/src/hydra_zen/wrapper/_implementations.py
@@ -3,14 +3,14 @@
 # pyright: strict, reportUnnecessaryTypeIgnoreComment = true, reportUnnecessaryIsInstance = false
 
 import warnings
-from collections import defaultdict, deque
+from collections import defaultdict
+from copy import deepcopy
 from functools import wraps
 from inspect import Parameter, signature
 from typing import (
     Any,
     Callable,
     DefaultDict,
-    Deque,
     Dict,
     FrozenSet,
     Generator,
@@ -1271,6 +1271,12 @@ class ZenStore:
     schema: <none>
     has_root: true
     _target_: __main__.Profile
+
+    **Manipulating and updating a store**
+
+    A store can be copied, updated, and merged. Its entries can have their groups
+    remapped, and individual entries can be deleted. See the docs for the corresponding
+    methods for details and examples.
     """
 
     __slots__ = (
@@ -1339,7 +1345,7 @@ class ZenStore:
         # created via the 'self-partialing' process
         self._internal_repo: Dict[Tuple[GroupName, NodeName], StoreEntry] = {}
         # Internal repo entries that have yet to be added to Hydra's config store
-        self._queue: Deque[StoreEntry] = deque([])
+        self._queue: Set[Tuple[GroupName, NodeName]] = set()
 
         self._deferred_to_config = deferred_to_config
         self._deferred_store = deferred_hydra_store
@@ -1409,8 +1415,7 @@ class ZenStore:
     def __call__(self: Self, __target: Optional[F] = None, **kw: Any) -> Union[F, Self]:
         """__call__(target : Optional[T] = None, /, name: NodeName | Callable[[Any], NodeName]] = ..., group: GroupName | Callable[[T], GroupName]] = None, package: Optional[str | Callable[[T], str]]] | None], provider: Optional[str], to_config: Callable[[T], Node] = ..., **to_config_kw: Any) -> T | ZenStore
 
-        The interface to an initialized store. Can be used to store a config or to
-        :ref:`customize the default values <self-partial>` of the store.
+        Store a config or :ref:`customize the default values <self-partial>` of the store.
 
         Parameters
         ----------
@@ -1544,20 +1549,118 @@ class ZenStore:
                 node=node,
             )
 
-            if not self._overwrite_ok and (_group, _name) in self._internal_repo:
-                raise ValueError(
-                    f"(name={entry['name']} group={entry['group']}): "
-                    f"Store entry already exists. Use a store initialized "
-                    f"with `ZenStore(overwrite_ok=True)` to overwrite config store "
-                    f"entries."
-                )
-            self._internal_repo[_group, _name] = entry
-            self._queue.append(entry)
-
-            if not self._deferred_store:
-                self.add_to_hydra_store()
+            self._set_entry(entry, overwrite=self._overwrite_ok)
             return cast(Union[F, Self], __target)
 
+    def copy(self: Self, store_name: Optional[str] = None) -> Self:
+        """Returns a copy of the store with the same overridden defaults.
+
+        Parameters
+        ----------
+        store_name : str | None, optional (default=None)
+
+        Returns
+        -------
+        ZenStore
+
+        Examples
+        --------
+        >>> from hydra_zen import ZenStore
+        >>> s1 = ZenStore()(group="G")
+        >>> s1({}, name="a")
+        >>> s2 = s1.copy()
+        >>> s2({}, name="b")
+        >>> s1
+        s1
+        {'G': ['a']}
+        >>> s2
+        s1_copy
+        {'G': ['a', 'b']}
+        """
+        cp = deepcopy(self)
+
+        cp.name = store_name if store_name is not None else self.name + "_copy"
+        return cp
+
+    def copy_with_mapped_groups(
+        self: Self,
+        old_group_to_new_group: Union[
+            Mapping[GroupName, GroupName], Callable[[GroupName], GroupName]
+        ],
+        *,
+        store_name: Optional[str] = None,
+        overwrite_ok: Optional[bool] = None,
+    ) -> Self:
+        """Create a copy of a store whose entries' groups have been updated according to the provided mapping.
+
+        Parameters
+        ----------
+        old_group_to_new_group : Mapping[GroupName, GroupName] | Callable[[GroupName], GroupName]
+            A mapping or callable that transforms an old group name to a new one.
+            Groups in the store that are not included in the mapping are unaffected.
+
+            A `GroupName` is `str | None`.
+
+        store_name : Optional[None]
+            If specified, the name of the new store.
+
+        overwrite_ok : Optional[bool]:
+            If specified, determines if the mapping can overwrite existing store
+            entries. Otherwise, defers to `ZenStore(overwrite_ok)`.
+
+        Returns
+        -------
+        new_store
+            A copy of `self` with remapped groups.
+
+        Examples
+        --------
+        >>> from hydra_zen import ZenStore
+
+        Creating an initial store
+
+        >>> s1 = ZenStore("s1")
+        >>> s1({}, group=None, name="a")
+        >>> s1({}, group="A/1", name="b")
+        >>> s1({}, group="A/2", name="c")
+        >>> s1
+        s1
+        {None: ['a'], 'A/1': ['b'], 'A/2': ['c']}
+
+        Replacing group "A/1" with "B", via a mapping
+
+        >>> s2 = s1.copy_with_mapped_groups({"A/1": "B"}, store_name="s2")
+        >>> s2
+        s2
+        {None: ['a'], 'A/2': ['c'], 'B': ['b']}
+
+        Placing all entries under group "A/" within a new inner group "p", via a
+        function
+
+        >>> s3 = s1.copy_with_mapped_groups(
+        ...     lambda g: g + "/p" if g and g.startswith("A/") else g, store_name="s3"
+        ... )
+        >>> s3
+        s3
+        {None: ['a'], 'A/1/p': ['b'], 'A/2/p': ['c']}
+        """
+        overwrite = overwrite_ok if overwrite_ok is not None else self._overwrite_ok
+
+        map_fn: Callable[[GroupName], GroupName] = (
+            (lambda x: old_group_to_new_group.get(x, x))
+            if isinstance(old_group_to_new_group, Mapping)
+            else old_group_to_new_group
+        )
+
+        copy = self.copy(store_name)
+        for (group, name), entry in tuple(copy._internal_repo.items()):
+            new_group = map_fn(group)
+            if new_group != group:
+                del copy[group, name]
+                entry["group"] = new_group
+                copy._set_entry(entry, overwrite=overwrite)
+        return copy
+
     @property
     def groups(self) -> Sequence[GroupName]:
         """Returns a sorted list of the groups registered with this store"""
@@ -1570,6 +1673,27 @@ class ZenStore:
             no_none = cast(Set[str], set_)
             return sorted(no_none)
 
+    def enqueue_all(self) -> None:
+        """Add all of the store's entries to the queue to be added to hydra's store.
+
+        Examples
+        --------
+        >>> from hydra_zen import ZenStore
+        >>> store = ZenStore(deferred_hydra_store=True)
+        >>> store({"a": 1}, name)
+        >>> store.has_enqueued()
+        True
+
+        >>> store.add_to_hydra_store()
+        >>> store.has_enqueued()
+        False
+
+        >>> store.enqueue_all()
+        >>> store.has_enqueued()
+        True
+        """
+        self._queue.update(self._internal_repo.keys())
+
     def has_enqueued(self) -> bool:
         """`True` if this store has entries that have not yet been added to
         Hydra's config store.
@@ -1600,6 +1724,84 @@ class ZenStore:
         not they have been added to Hydra's config store"""
         return bool(self._internal_repo)
 
+    def __len__(self) -> int:
+        return len(self._internal_repo)
+
+    def update(self, __other: "ZenStore") -> None:
+        """Updates the store inplace with redundant entries being overwritten.
+
+        Can also be applied via the `|=` in-place operator.
+
+        Examples
+        --------
+        >>> from hydra_zen import ZenStore
+        >>> def f(): ...
+        >>> def g(): ...
+        >>> s1 = ZenStore("s1")
+        >>> s2 = ZenStore("s2")
+        >>> s1(f)  # store f in s1
+        >>> s2(g)  # store g in s2
+
+        >>> s1.update(s2)
+        >>> s1  # s1 now has entries for both f and g
+        s1
+        {None: ['f', 'g']}
+
+        Alternatively, the `|=` operator can be used to update a store inplace.
+
+        >>> s3 = ZenStore("s3")
+        >>> s3 |= s2
+        >>> s3
+        s3
+        {None: ['g']}
+        """
+        self._internal_repo.update(deepcopy(__other._internal_repo))
+        self._queue.update(__other._queue)
+        if not self._deferred_store:
+            self.add_to_hydra_store()
+        return
+
+    def merge(
+        self: Self, __other: "ZenStore", store_name: Optional[str] = None
+    ) -> Self:
+        """Create a new store by merging two stores.
+
+        The new store's default settings will reflect those of `self` in
+        `self.merge(other)`. This can also be applied via the `|` operator.
+
+        Examples
+        --------
+        >>> from hydra_zen import ZenStore
+        >>> def f(): ...
+        >>> def g(): ...
+        >>> s1 = ZenStore("s1")
+        >>> s2 = ZenStore("s2")
+        >>> s1(f)  # store f in s1
+        >>> s2(g)  # store g in s2
+
+        >>> s3 = s1.merge(s2)
+        >>> s3
+        s1_copy
+        {None: ['f', 'g']}
+
+        Alternatively, the `|` operator can be used to merge stores.
+
+        >>> s4 = s1 | s2
+        >>> s4
+        s1_copy
+        {None: ['f', 'g']}
+        """
+        cp = self.copy(store_name)
+        cp.update(__other)
+        return cp
+
+    def __or__(self: Self, other: "ZenStore") -> Self:
+        return self.merge(other)
+
+    def __ior__(self: Self, other: "ZenStore") -> Self:
+        self.update(other)
+        return self
+
     @overload
     def __getitem__(self, key: Tuple[GroupName, NodeName]) -> Node:
         ...
@@ -1661,10 +1863,31 @@ class ZenStore:
             }
         return _resolve_node(self._internal_repo[key], copy=False)["node"]
 
+    def __delitem__(self, key: Tuple[GroupName, NodeName]) -> None:
+        del self._internal_repo[key]
+        self._queue.discard(key)
+
+    def delete_entry(self, group: GroupName, name: NodeName) -> None:
+        del self[group, name]
+
     def get_entry(self, group: GroupName, name: NodeName) -> StoreEntry:
         """Access a store entry, which is a mapping that specifies the entry's
         name, group, package, provider, and node.
 
+        Parameters
+        ----------
+        group : str | None
+        name : str
+
+        Returns
+        -------
+        dict
+           - name: NodeName
+           - group: GroupName
+           - package: Optional[str]
+           - provider: Optional[str]
+           - node: ConfigType
+
         Notes
         -----
         Mutating the returned mapping will not affect the store's internal entry.
@@ -1684,6 +1907,22 @@ class ZenStore:
         """
         return _resolve_node(self._internal_repo[(group, name)], copy=True)
 
+    def _set_entry(self, __entry: StoreEntry, overwrite: bool) -> None:
+        _group = __entry["group"]
+        _name = __entry["name"]
+        if not overwrite and (_group, _name) in self._internal_repo:
+            raise ValueError(
+                f"(name={__entry['name']} group={__entry['group']}): "
+                f"Store entry already exists. Use a store initialized "
+                f"with `ZenStore(overwrite_ok=True)` to overwrite config store "
+                f"entries."
+            )
+        self._internal_repo[_group, _name] = __entry
+        self._queue.add((_group, _name))
+
+        if not self._deferred_store:
+            self.add_to_hydra_store()
+
     def __contains__(self, key: Union[GroupName, Tuple[GroupName, NodeName]]) -> bool:
         """Checks if group or (group, node-name) exists in zen-store."""
         if key is None:
@@ -1760,8 +1999,9 @@ class ZenStore:
 
         """
         _store = ConfigStore.instance().store
-        while self._queue:
-            entry = _resolve_node(self._queue.popleft(), copy=False)
+
+        for key in tuple(self._queue):
+            entry = _resolve_node(self._internal_repo[key], copy=False)
             if (
                 (
                     overwrite_ok is False
@@ -1783,6 +2023,7 @@ class ZenStore:
                     f"`overwrite_ok=True` to enable replacing config store entries"
                 )
             _store(**entry)
+            self._queue.discard(key)
 
     def _exists_in_hydra_store(
         self,


Test patch:
diff --git a/tests/test_dataclass_semantics.py b/tests/test_dataclass_semantics.py
index 82755be0..7061062b 100644
--- a/tests/test_dataclass_semantics.py
+++ b/tests/test_dataclass_semantics.py
@@ -17,6 +17,7 @@ from hydra_zen import (
     just,
     make_config,
     make_custom_builds_fn,
+    to_yaml,
 )
 from hydra_zen.errors import HydraZenDeprecationWarning
 
@@ -451,3 +452,14 @@ def test_module_is_not_specified(Conf):
 )
 def test_modules_is_specified(Conf):
     assert Conf.__module__ == "aaa"
+
+
+@pytest.mark.parametrize(
+    "cfg",
+    [
+        builds(dict, x=1),
+        builds(dict, x=2)(),
+    ],
+)
+def test_builds_is_copyable(cfg):
+    assert to_yaml(cfg) == to_yaml(deepcopy(cfg))
diff --git a/tests/test_store.py b/tests/test_store.py
index 37cd115d..2cebfe9b 100644
--- a/tests/test_store.py
+++ b/tests/test_store.py
@@ -426,6 +426,7 @@ def test_overwrite_ok(outer: bool, inner: bool, name, group):
         ):
             _store({}, name=name, group=group)
         return
+    _store.add_to_hydra_store(overwrite_ok=True)
     _store({}, name=name, group=group)
     if not inner:
         with pytest.raises(
@@ -548,7 +549,7 @@ def test_default_to_config_validates_dataclass_instance_with_kw():
 )
 def test_validate_init(param_name):
     with pytest.raises(TypeError, match=f"{param_name} must be a bool"):
-        ZenStore(**{param_name: "bad"})
+        ZenStore(**{param_name: "bad"})  # type: ignore
 
 
 def contains(key, store_):
@@ -808,3 +809,173 @@ def test_store_hydrated_dataclass():
     store = ZenStore()
     store(SomethingHydrated, name="foo", x=2)
     assert instantiate(store[None, "foo"]) == 2
+
+
+def test_del():
+    s = ZenStore()
+    s({}, name="a")
+    s({}, name="b")
+    s2 = s(group="G")
+    assert len(s) == 2
+    assert len(s._queue) == 2
+
+    del s[None, "a"]
+    assert len(s) == 1
+    assert len(s._queue) == 1
+    assert (None, "a") not in s
+    assert (None, "a") not in s._queue
+
+    assert s
+    assert s.has_enqueued()
+    assert s == s2
+
+    s.delete_entry(None, "b")
+    assert not s
+    assert not s.has_enqueued()
+
+
+def test_copy():
+    s = ZenStore(name="s")(group="G")
+    s({}, name="a")
+    s({}, name="b")
+    s2 = s(group="G")
+    s3 = s.copy()
+    s3({}, name="c")
+
+    assert s == s2
+    assert s != s3
+
+    del s["G", "a"]
+    assert len(s) == 1
+    assert len(s._queue) == 1
+    assert len(s3) == 3
+    assert len(s3._queue) == 3
+    assert ("G", "c") in s3
+
+    assert s3.name == "s_copy"
+    assert s3.copy("moo").name == "moo"
+
+
+@pytest.mark.usefixtures("clean_store")
+@pytest.mark.parametrize(
+    "mapping",
+    [
+        {"A": "B"},
+        lambda x: "B" if x and x.startswith("A") else x,
+    ],
+)
+def test_map_groups(mapping):
+    s1 = ZenStore()
+    s1({"x": 1}, group=None, name="a")
+    s1({"y": 2}, group="A", name="b")
+    s1({"z": 3}, group="A", name="c")
+
+    s2 = s1.copy_with_mapped_groups(mapping)
+    assert s1 != s2
+
+    assert (None, "a") in s1
+    assert ("A", "b") in s1
+    assert ("B", "b") not in s1
+    assert len(s1) == 3
+    assert s1._queue == {(None, "a"), ("A", "b"), ("A", "c")}
+
+    assert (None, "a") in s1
+    assert ("A", "b") not in s2
+    assert ("B", "b") in s2
+    assert len(s2) == 3
+    assert s2._queue == {(None, "a"), ("B", "b"), ("B", "c")}
+
+    s1.add_to_hydra_store()
+    assert instantiate_from_repo("a", group=None) == {"x": 1}
+    assert instantiate_from_repo("b", group="A") == {"y": 2}
+
+    with pytest.raises(KeyError):
+        instantiate_from_repo("b", group="B")
+
+    s2.add_to_hydra_store(overwrite_ok=True)
+    assert instantiate_from_repo("b", group="B") == {"y": 2}
+
+
+@pytest.mark.usefixtures("clean_store")
+def test_map_with_overwrite():
+    s1 = ZenStore()
+    s1(dict(x=1), name="b")
+    s1(dict(x=2), name="b", group="G")
+    with pytest.raises(ValueError, match="Store entry already exists"):
+        s1.copy_with_mapped_groups({None: "G"})
+
+    s2 = s1.copy_with_mapped_groups({None: "G"}, overwrite_ok=True)
+    s2.add_to_hydra_store()
+    assert instantiate_from_repo("b", group="G") == {"x": 1}
+
+
+@pytest.mark.usefixtures("clean_store")
+def test_enqueue_all():
+    s1 = ZenStore()
+    s1(dict(x=1), name="b")
+    s2 = s1(group="G")
+    s2(dict(x=2), name="b")
+
+    q = s2._queue.copy()
+    assert s1 == s2
+    s2.add_to_hydra_store(overwrite_ok=True)
+    assert s1 == s2
+    assert not s1.has_enqueued()
+    s1.enqueue_all()
+    assert s1.has_enqueued()
+    assert s1._queue == q
+    assert s1 == s2
+
+
+def test_update():
+    s1 = ZenStore()
+    s2 = ZenStore()
+    s1({}, name="a")
+    s1({}, name="b")  # should get overwritten
+
+    s2({"x": 1}, name="b")
+    s2({}, name="c", group="G")  # should be added
+
+    s3 = s1(group="BB")
+    assert s3 == s1
+    s1 |= s2
+
+    assert len(s1) == 3
+    assert s1._queue == {(None, "a"), (None, "b"), ("G", "c")}
+    assert s1._internal_repo[None, "b"] is not s2._internal_repo[None, "b"]
+    assert s2[None, "b"] == {"x": 1}
+
+    # make sure partiald stores are still connected
+    assert s3 == s1
+    assert ("BB", "foo") not in s1
+    s3({}, name="foo")
+    assert ("BB", "foo") in s1
+
+
+@pytest.mark.usefixtures("clean_store")
+def test_update_respects_add_to_hydra_store():
+    s = ZenStore(deferred_hydra_store=True)
+    s({}, name="a")
+    s1 = ZenStore(deferred_hydra_store=True)
+    s2 = ZenStore(deferred_hydra_store=False)
+    s1.update(s)
+    with pytest.raises(KeyError):
+        instantiate_from_repo("a")
+
+    s2.update(s)
+    assert instantiate_from_repo("a") == {}
+
+
+def test_merge():
+    s1 = ZenStore()
+    s2 = ZenStore()
+    s1({}, name="a")
+    s2({}, name="b")
+    s3 = s1 | s2
+    s3({}, name="c")
+    assert len(s1) == 1
+    assert len(s2) == 1
+    assert len(s3) == 3
+    assert s3._queue == {(None, "a"), (None, "b"), (None, "c")}
+    assert s3._internal_repo[None, "a"] is not s1._internal_repo[None, "a"]
+    assert s3._internal_repo[None, "b"] is not s2._internal_repo[None, "b"]