Repository: mit-ll-responsible-ai/hydra-zen

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? 


Discussion:
@petosa I'll think about it (and I'm sure @rsokl will have more insight) but I don't think this is possible in hydra-zen as the resolver is handled by Hydra.  You might want to look into Hydra callbacks: https://hydra.cc/docs/experimental/callbacks/
Thanks for your reply. Couldn't the requested change be achieved by moving line 336 in the linked code to line 340?
Yeah, I can try and look at it.  @rsokl might have a better understanding of the dynamics at play between hydra-zen and Hydra here though, but we may not want to mess with that behavior.  It still might be better to use callbacks (which can be configured with hydra-zen).   I'm not sure when the resolver is called for callbacks though.  Do you have a simple example of what you are trying to do? Maybe there's another solution?
I like the idea of being able to use `pre_call` to do stuff to the cfg before it gets resolved. Moving `resolve` is no good because it could easily break people's existing code.

Off the top of my head, I think `Zen`/`zen` should have a `resolve_cfg: bool = True` param. For your desired case, you would pass `resolve_cfg=False` then have your `pre_call` function be responsible for calling `OmegaConf.resolve` at the end. This would permit maximum flexibility.

I would happily accept a PR for this. Our CI seems to need some tending to due to breaking changes in pydantic and some other flakiness, so be warned that it might be a little noisy. I hope to fix it this weekend.
(CI is back in good health)

Resolving patch:
diff --git a/src/hydra_zen/wrapper/_implementations.py b/src/hydra_zen/wrapper/_implementations.py
--- 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 @@ def __init__(
         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 @@ def __init__(
                 "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 @@ def __init__(
             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 @@ def __call__(self, __cfg: Union[ConfigLike, str]) -> 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 task(x: int, y: int):
     """
     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 @@ def __init__(
             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
--- 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