Repository: mit-ll-responsible-ai/hydra-zen
Pull request #236 (2022-02-26 18:41:06)

Fix bug: parent configs can introduce zen-processing via inheritance
This PR fixes a bug where inheriting from configs that leverage zen-processing features was not handled correctly. 

This PR also prevents users from manually inheriting the `_partial_` field in configs that use zen-processing via `builds`; this scenario will raise a `TypeError`.

The buggy behavior is exemplified by the following scenario:

**Before**:

```python
>>> from hydra_zen import builds, instantiate, to_yaml

>>> Parent = builds(dict, a="${.s}", zen_meta=dict(s=1))
>>> instantiate(Parent)
{'a': 1}

>>> Child = builds(dict, builds_bases=(Parent,))
>>> instantiate(Child)
{'_zen_target': 'builtins.dict', '_zen_exclude': ['s'], 's': 1, 'a': 1}
```

`Parent` and `Child` ought to instantiate to the same result: `{'a': 1}`, however this does not occur because `builds` did not handle inheritance properly; see that `_target_` gets overwritten via naive inheritance:

```python
>>> print(to_yaml(Parent))
_target_: hydra_zen.funcs.zen_processing
_zen_target: builtins.dict
_zen_exclude:
- s
a: ${.s}
s: 1

>>> print(to_yaml(Child))
_target_: builtins.dict
_zen_target: builtins.dict
_zen_exclude:
- s
a: ${.s}
s: 1
```

**After**:

`builds` will now detect the usage of zen-processing among the parent configs in `builds_bases`, and will dispatch to `zen_processing` in the resulting config as-appropriate. E.g.

```python
>>> from hydra_zen import builds, instantiate, to_yaml

>>> Parent = builds(dict, a="${.s}", zen_meta=dict(s=1))
>>> instantiate(Parent)
{'a': 1}

>>> Child = builds(dict, builds_bases=(Parent,))
>>> instantiate(Child)
{'a': 1}
```






Patch:
diff --git a/docs/source/changes.rst b/docs/source/changes.rst
index 3b40e88b..6d90b1d1 100644
--- a/docs/source/changes.rst
+++ b/docs/source/changes.rst
@@ -11,7 +11,7 @@ chronological order. All previous releases should still be available on pip.
 .. _v0.6.0:
 
 ---------------------
-0.6.0rc2 - 2022-02-17
+0.6.0rc3 - 2022-02-26
 ---------------------
 
 This release focuses on improving hydra-zen's type-annotations; it increases the 
@@ -25,6 +25,11 @@ IDEs (e.g. using Pylance in VSCode) for the best experience.
 
 (A note to VSCode users: make sure to set `Type Checking Mode` to `basic` in your IDE -- it is disabled by default!)
 
+Bug Fixes
+---------
+
+``builds(<target>, builds_bases=(...))`` now properly supports the case where a parent config introduces zen-processing features via inheritance. See pull:`236` for more details.
+
 
 Improvements
 ------------
diff --git a/src/hydra_zen/structured_configs/_implementations.py b/src/hydra_zen/structured_configs/_implementations.py
index 31287546..82e8965e 100644
--- a/src/hydra_zen/structured_configs/_implementations.py
+++ b/src/hydra_zen/structured_configs/_implementations.py
@@ -733,9 +733,7 @@ def builds(
     Type[PartialBuilds[Importable]],
     Type[BuildsWithSig[Type[R], P]],
 ]:
-    """builds(hydra_target, /, *pos_args, zen_partial=False, zen_meta=None,
-    hydra_recursive=None, populate_full_signature=False, hydra_convert=None,
-    frozen=False, dataclass_name=None, builds_bases=(), **kwargs_for_target)
+    """builds(hydra_target, /, *pos_args, zen_partial=False, zen_wrappers=(), zen_meta=None, populate_full_signature=False, hydra_recursive=None, hydra_convert=None, frozen=False, dataclass_name=None, builds_bases=(), **kwargs_for_target)
 
     Returns a structured config, which describes how to instantiate/call
     ``<hydra_target>`` with both user-specified and auto-populated parameter values.
@@ -1208,13 +1206,14 @@ def builds(
 
     target_path: Final[str] = _utils.get_obj_path(target)
 
-    if (
-        HYDRA_SUPPORTS_PARTIAL
-        and zen_partial
-        # check that no other zen-processing is needed
-        and not zen_meta
-        and not validated_wrappers
-    ):  # pragma: no cover
+    requires_zen_processing: Final[bool] = (
+        bool(zen_meta)
+        or bool(validated_wrappers)
+        or any(uses_zen_processing(b) for b in builds_bases)
+        or (zen_partial and not HYDRA_SUPPORTS_PARTIAL)
+    )
+
+    if not requires_zen_processing and zen_partial:  # pragma: no cover
         # TODO: require test-coverage once Hydra publishes nightly builds
         target_field = [
             (
@@ -1228,7 +1227,7 @@ def builds(
                 _utils.field(default=zen_partial, init=False),
             ),
         ]
-    elif zen_partial or zen_meta or validated_wrappers:
+    elif requires_zen_processing:
         # target is `hydra_zen.funcs.zen_processing`
         target_field = [
             (
@@ -1715,20 +1714,26 @@ def builds(
         dataclass_name, fields=sanitized_base_fields, bases=builds_bases, frozen=frozen
     )
 
-    # TODO: revisit this constraint
     if zen_partial is False and (
-        hasattr(out, _ZEN_PARTIAL_TARGET_FIELD_NAME)
-        or (HYDRA_SUPPORTS_PARTIAL and hasattr(out, _PARTIAL_FIELD_NAME))
+        getattr(out, _ZEN_PARTIAL_TARGET_FIELD_NAME, False)
+        or getattr(out, _PARTIAL_FIELD_NAME, False)
     ):
-        # `out._partial_` has been inherited; this will lead to an error when
-        # hydra-instantiation occurs, since it will be passed to target.
-        # There is not an easy way to delete this, since it comes from a parent class
+        # `out._partial_=True` or `out._zen_partial=True` has been inherited; there
+        # is no way for users to override this, thus they must explicitly specify
+        # `zen_partial=True` in order to "opt-in" to this behavior.
         raise TypeError(
             BUILDS_ERROR_PREFIX
             + "`builds(..., zen_partial=False, builds_bases=(...))` does not "
             "permit `builds_bases` where a partial target has been specified."
         )
 
+    if requires_zen_processing and hasattr(out, _PARTIAL_FIELD_NAME):
+        raise TypeError(
+            BUILDS_ERROR_PREFIX
+            + "`builds(..., builds_bases=(...))` cannot use zen-processing features "
+            "while inheriting the field `_partial_: bool = ...`"
+        )
+
     out.__doc__ = (
         f"A structured config designed to {'partially ' if zen_partial else ''}initialize/call "
         f"`{target_path}` upon instantiation by hydra."
@@ -1739,6 +1744,9 @@ def builds(
             out.__doc__ += (
                 f"\n\nThe docstring for {_utils.safe_name(target)} :\n\n" + target_doc
             )
+
+    assert requires_zen_processing is uses_zen_processing(out)
+
     return cast(Union[Type[Builds[Importable]], Type[BuildsWithSig[Type[R], P]]], out)
 
 


Test patch:
diff --git a/tests/__init__.py b/tests/__init__.py
index 88320547..73a0738e 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -1,8 +1,11 @@
 # Copyright (c) 2022 Massachusetts Institute of Technology
 # SPDX-License-Identifier: MIT
 
+from typing import Any
+
 import hypothesis.strategies as st
 
+from hydra_zen import to_yaml
 from hydra_zen.structured_configs._utils import is_classmethod
 
 valid_hydra_literals = st.sampled_from(
@@ -23,3 +26,7 @@ def check_identity(new, original):
         return new is original
     # objects like classmethods do not satisfy `x is x`
     return new == original
+
+
+def sorted_yaml(conf: Any):
+    return "\n".join(sorted(to_yaml(conf).splitlines()))
diff --git a/tests/test_zen_processing/test_inheritance.py b/tests/test_zen_processing/test_inheritance.py
new file mode 100644
index 00000000..d8c42834
--- /dev/null
+++ b/tests/test_zen_processing/test_inheritance.py
@@ -0,0 +1,155 @@
+from dataclasses import dataclass
+from functools import partial
+
+import hypothesis.strategies as st
+import pytest
+from hypothesis import given
+
+from hydra_zen import builds, instantiate, to_yaml
+from hydra_zen._compatibility import HYDRA_SUPPORTS_PARTIAL
+from hydra_zen.structured_configs._implementations import (
+    _PARTIAL_FIELD_NAME,
+    _ZEN_PROCESSING_LOCATION,
+    uses_zen_processing,
+)
+from hydra_zen.typing import HydraPartialBuilds, ZenPartialBuilds
+from tests import sorted_yaml
+
+
+def wrapper1(func):
+    return func
+
+
+def wrapper2(func):
+    return func
+
+
+@given(
+    args=st.dictionaries(st.sampled_from("abc"), st.integers(-5, 5)),
+    meta=st.none() | st.dictionaries(st.sampled_from("def"), st.integers(-5, 5)),
+    wrappers=st.lists(st.sampled_from([None, wrapper1, wrapper2])),
+    meta_via_parent=st.booleans(),
+    wrappers_via_parent=st.booleans(),
+)
+def test_inherited_zen_processing(
+    args, meta, wrappers, meta_via_parent, wrappers_via_parent
+):
+    parent_kwargs = {}
+    child_kwargs = {}
+
+    target = parent_kwargs if meta_via_parent else child_kwargs
+    target["zen_meta"] = meta
+
+    target = parent_kwargs if wrappers_via_parent else child_kwargs
+    target["zen_wrappers"] = wrappers
+
+    parent = builds(dict, **args, **parent_kwargs)
+    child = builds(dict, **child_kwargs, builds_bases=(parent,))
+
+    direct = builds(dict, **args, zen_meta=meta, zen_wrappers=wrappers)
+
+    assert sorted_yaml(direct) == sorted_yaml(child)
+    assert instantiate(direct) == instantiate(child)
+
+
+def yaml_lines(conf):
+    return set(to_yaml(conf).splitlines())
+
+
+@given(
+    args=st.dictionaries(st.sampled_from("abc"), st.integers(-5, 5)),
+    meta=st.none() | st.dictionaries(st.sampled_from("def"), st.integers(-5, 5)),
+    wrappers=st.lists(st.sampled_from([None, wrapper1, wrapper2])),
+)
+def test_partial_with_inherited_zen_processing(args, meta, wrappers):
+    parent = builds(dict, **args, zen_meta=meta, zen_wrappers=wrappers)
+    Child = builds(dict, zen_partial=True, builds_bases=(parent,))
+
+    child = Child()
+
+    if HYDRA_SUPPORTS_PARTIAL or uses_zen_processing(parent):
+        # ensure meta fields are retained by child
+        assert yaml_lines(parent) <= yaml_lines(child)
+
+    if uses_zen_processing(parent) or not HYDRA_SUPPORTS_PARTIAL:
+        assert uses_zen_processing(child)
+        assert isinstance(child, ZenPartialBuilds)
+        assert not isinstance(child, HydraPartialBuilds)
+    else:
+        assert not uses_zen_processing(child)
+        assert isinstance(child, HydraPartialBuilds)
+
+    out = instantiate(child)
+    assert isinstance(out, partial)
+    assert out() == args
+
+
+@dataclass
+class ZenPartialFalse:
+    _target_: str = _ZEN_PROCESSING_LOCATION
+    _zen_target: str = "builtins.dict"
+    _zen_partial: bool = False
+    a: int = 1
+
+
+@dataclass
+class ZenPartialTrue:
+    _target_: str = _ZEN_PROCESSING_LOCATION
+    _zen_target: str = "builtins.dict"
+    _zen_partial: bool = True
+    a: int = 1
+
+
+@dataclass
+class HydraPartialFalse:
+    _target_: str = "builtins.dict"
+    _partial_: bool = False
+    a: int = 1
+
+
+@dataclass
+class HydraPartialTrue:
+    _target_: str = "builtins.dict"
+    _partial_: bool = True
+    a: int = 1
+
+
+@pytest.mark.parametrize("Parent", [ZenPartialFalse, HydraPartialFalse])
+def test_inherited_partial_false_fields(Parent):
+    if not HYDRA_SUPPORTS_PARTIAL and Parent is HydraPartialFalse:
+        pytest.mark.skip("Hydra version doesn't support instantiate API")
+
+    Conf = builds(dict, builds_bases=(Parent,))  # should be OK
+    assert instantiate(Conf) == instantiate(Parent)
+
+
+def test_inherit_hydra_partial_with_zen_processing_raises():
+    with pytest.raises(TypeError):
+        builds(
+            dict, zen_meta=dict(a=1), zen_partial=True, builds_bases=(HydraPartialTrue,)
+        )
+
+
+@pytest.mark.parametrize("Parent", [ZenPartialTrue, HydraPartialTrue])
+def test_zen_partial_false_raises_for_partiald_parents(Parent):
+    with pytest.raises(TypeError):
+        builds(dict, zen_partial=False, builds_bases=(Parent,))
+
+
+@pytest.mark.parametrize(
+    "Parent", [ZenPartialFalse, HydraPartialFalse, ZenPartialTrue, HydraPartialTrue]
+)
+def test_zen_partial_true_holds_for_all_inheritance(Parent):
+    def make_conf():
+        return builds(dict, b=2, zen_partial=True, builds_bases=(Parent,))
+
+    if not HYDRA_SUPPORTS_PARTIAL and hasattr(Parent, _PARTIAL_FIELD_NAME):
+        with pytest.raises(TypeError):
+            make_conf()
+        return
+
+    Conf = make_conf()
+
+    out = instantiate(Conf)
+    assert isinstance(out, partial)
+    assert out() == {"a": 1, "b": 2}