Repository: mit-ll-responsible-ai/hydra-zen (python)

Issue:
Move `OmegaConf.resolve(cfg)` after `pre_call`
Hello,

I am using the `@zen` decorator in my project, and using its `pre_call` arg to modify my config. I am trying to modify the unresolved config, but `pre_call` gets the config after `OmegaConf.resolve(cfg)` has been run.
https://github.com/mit-ll-responsible-ai/hydra-zen/blob/796d8ac7c87c5ad245eb739cd27b81cf62eb19ca/src/hydra_zen/wrapper/_implementations.py#L336

Can `OmegaConf.resolve(cfg)` please be moved to be after `self.pre_call(cfg)` to allow for more flexibility? 

Pull request:
make config resolution configurable in zen
Closes #444 

Patch:
diff --git a/docs/source/changes.rst b/docs/source/changes.rst
index 56ef9293..ece432cc 100644
--- a/docs/source/changes.rst
+++ b/docs/source/changes.rst
@@ -56,6 +56,7 @@ Improvements
 ------------
 - :func:`~hydra_zen.builds` now has a transitive property that enables iterative build patterns. See :pull:`455`
 - :func:`~hydra_zen.zen`'s instantiation phase has been improved so that dataclass objects and stdlib containers are returned instead of omegaconf objects. See :pull:`#448`. 
+- :func:`~hydra_zen.zen` can now be passed `resolve_pre_call=False` to defer the resolution of interpolated fields until after `pre_call` functions are called. See :pull:`460`.
 
 Bug Fixes
 ---------
diff --git a/src/hydra_zen/wrapper/_implementations.py b/src/hydra_zen/wrapper/_implementations.py
index 6fcc9144..3301979e 100644
--- a/src/hydra_zen/wrapper/_implementations.py
+++ b/src/hydra_zen/wrapper/_implementations.py
@@ -1,6 +1,6 @@
 # Copyright (c) 2023 Massachusetts Institute of Technology
 # SPDX-License-Identifier: MIT
-# pyright: strict
+# pyright: strict, reportUnnecessaryTypeIgnoreComment = true, reportUnnecessaryIsInstance = false
 
 import warnings
 from collections import defaultdict, deque
@@ -136,6 +136,7 @@ class Zen(Generic[P, R]):
         exclude: Optional[Union[str, Iterable[str]]] = None,
         pre_call: PreCall = None,
         unpack_kwargs: bool = False,
+        resolve_pre_call: bool = True,
     ) -> None:
         """
         Parameters
@@ -173,9 +174,14 @@ class Zen(Generic[P, R]):
                 "signatures."
             )
 
-        if not isinstance(unpack_kwargs, bool):  # type: ignore
+        if not isinstance(unpack_kwargs, bool):
             raise TypeError(f"`unpack_kwargs` must be type `bool` got {unpack_kwargs}")
 
+        if not isinstance(resolve_pre_call, bool):  # pragma: no cover
+            raise TypeError(
+                f"`resolve_pre_call` must be type `bool` got {resolve_pre_call}"
+            )
+        self._resolve = resolve_pre_call
         self._unpack_kwargs: bool = unpack_kwargs and any(
             p.kind is p.VAR_KEYWORD for p in self.parameters.values()
         )
@@ -222,8 +228,8 @@ class Zen(Generic[P, R]):
             pre_call if not isinstance(pre_call, Iterable) else _flat_call(pre_call)
         )
 
-    @staticmethod
     def _normalize_cfg(
+        self,
         cfg: Union[
             DataClass_,
             Type[DataClass_],
@@ -340,8 +346,10 @@ class Zen(Generic[P, R]):
             The result of `func(<args extracted from cfg>)`
         """
         cfg = self._normalize_cfg(__cfg)
-        # resolves all interpolated values in-place
-        OmegaConf.resolve(cfg)
+
+        if self._resolve:
+            # resolves all interpolated values in-place
+            OmegaConf.resolve(cfg)
 
         if self.pre_call is not None:
             self.pre_call(cfg)
@@ -461,6 +469,7 @@ def zen(
     unpack_kwargs: bool = ...,
     pre_call: PreCall = ...,
     ZenWrapper: Type[Zen[P, R]] = Zen,
+    resolve_pre_call: bool = ...,
     exclude: Optional[Union[str, Iterable[str]]] = None,
 ) -> Zen[P, R]:
     ...
@@ -472,6 +481,7 @@ def zen(
     *,
     unpack_kwargs: bool = ...,
     pre_call: PreCall = ...,
+    resolve_pre_call: bool = ...,
     ZenWrapper: Type[Zen[Any, Any]] = ...,
     exclude: Optional[Union[str, Iterable[str]]] = None,
 ) -> Callable[[Callable[P, R]], Zen[P, R]]:
@@ -484,6 +494,7 @@ def zen(
     unpack_kwargs: bool = False,
     pre_call: PreCall = None,
     exclude: Optional[Union[str, Iterable[str]]] = None,
+    resolve_pre_call: bool = True,
     ZenWrapper: Type[Zen[P, R]] = Zen,
 ) -> Union[Zen[P, R], Callable[[Callable[P, R]], Zen[P, R]]]:
     r"""zen(func, /, pre_call, ZenWrapper)
@@ -521,12 +532,16 @@ def zen(
         This is useful, e.g., for seeding a RNG prior to the instantiation phase
         that is triggered when calling the wrapped function.
 
-    exclude: Optional[str | Iterable[str]]
+    exclude : Optional[str | Iterable[str]]
         Specifies one or more parameter names in the function's signature
         that will not be extracted from input configs by the zen-wrapped function.
 
         A single string of comma-separated names can be specified.
 
+    resolve_pre_call : bool, (default=True)
+        If True, the config passed to the zen-wrapped function has its interpolated
+        fields resolved to being passed to any pre-call functions.
+
     ZenWrapper : Type[hydra_zen.wrapper.Zen], optional (default=Zen)
         If specified, a subclass of `Zen` that customizes the behavior of the wrapper.
 
@@ -722,12 +737,20 @@ def zen(
     """
     if __func is not None:
         return ZenWrapper(
-            __func, pre_call=pre_call, exclude=exclude, unpack_kwargs=unpack_kwargs
+            __func,
+            pre_call=pre_call,
+            exclude=exclude,
+            unpack_kwargs=unpack_kwargs,
+            resolve_pre_call=resolve_pre_call,
         )
 
     def wrap(f: Callable[P, R]) -> Zen[P, R]:
         return ZenWrapper(
-            f, pre_call=pre_call, exclude=exclude, unpack_kwargs=unpack_kwargs
+            f,
+            pre_call=pre_call,
+            exclude=exclude,
+            unpack_kwargs=unpack_kwargs,
+            resolve_pre_call=resolve_pre_call,
         )
 
     return wrap
@@ -1278,15 +1301,15 @@ class ZenStore:
             with `store(node=Config)` instead of actually storing the node with
             `store(Config)`.
         """
-        if not isinstance(deferred_to_config, bool):  # type: ignore
+        if not isinstance(deferred_to_config, bool):
             raise TypeError(
                 f"deferred_to_config must be a bool, got {deferred_to_config}"
             )
 
-        if not isinstance(overwrite_ok, bool):  # type: ignore
+        if not isinstance(overwrite_ok, bool):
             raise TypeError(f"overwrite_ok must be a bool, got {overwrite_ok}")
 
-        if not isinstance(deferred_hydra_store, bool):  # type: ignore
+        if not isinstance(deferred_hydra_store, bool):
             raise TypeError(
                 f"deferred_hydra_store must be a bool, got {deferred_hydra_store}"
             )


Test patch:
diff --git a/tests/test_zen.py b/tests/test_zen.py
index 8fffa0b3..b0b32c10 100644
--- a/tests/test_zen.py
+++ b/tests/test_zen.py
@@ -217,6 +217,27 @@ def test_zen_resolves_default_factories():
     assert zen_identity(Cfg()) == [1, 2, 3]
 
 
+def test_no_resolve():
+    def not_resolved(x):
+        assert x == dict(x=1, y="${x}")
+
+    def is_resolved(x):
+        assert x == dict(x=1, y=1)
+
+    Cfg = make_config(x=1, y="${x}")
+    out = zen(lambda **kw: kw, unpack_kwargs=True, pre_call=is_resolved)(Cfg)
+    assert out == dict(x=1, y=1)
+
+    Cfg2 = make_config(x=1, y="${x}")
+    out2 = zen(
+        lambda **kw: kw,
+        unpack_kwargs=True,
+        resolve_pre_call=False,
+        pre_call=not_resolved,
+    )(Cfg2)
+    assert out2 == dict(x=1, y=1)
+
+
 def test_zen_works_on_partiald_funcs():
     from functools import partial