Implements zen-wrappers

[rsokl]

This adds a new feature: builds( .... zen_wrappers=<..>).

This is an extremely powerful feature that enables one to inject modifications to the instantiation process by including wrappers in a target's configuration.

The Basics

Whereas

Conf = builds(f, *args, **kwargs)
instantiate(Conf)

calls f(*args, **kwargs), the following config with wrappers

Conf = builds(f, *args, **kwargs, zen_target=[w1, w2, w3])

will instantiate as:

f = w1(f)
f = w2(f)
f = w3(f)
f(*args, **kwargs)

The wrapper paradigm enables arbitrary pre-processing, post-processing, and transformation of the target and its inputs/outputs.

Seeing it in Action

Suppose we have a function that takes in a temperature in Farenheit, and converts it to Celcius.

def faren_to_celsius(temp_f):
    return ((temp_f - 32) * 5) / 9

>>> AsCelcius = builds(faren_to_celsius)
>>> instantiate(AsCelcius, temp_f=32)
0.0

But we discover that we need the temperature to be converted to Kelvin instead. We can modify our config to include the following wrapper:

>>> def change_celcius_to_kelvin(celc_func):
...     def wraps(*args, **kwargs): 
...         return 273.15 + celc_func(*args, **kwargs) 
...     return wraps

>>> AsKelvin = builds(faren_to_celsius, zen_wrappers=change_celcius_to_kelvin)
>>> instantiate(AsKelvin, temp_f=32)
273.15

Another use case

In the next PR, users will be able to add pydantic validation to the instantiation process.

from pydantic import PositiveInt
from hydra_zen.experimental.third_party.pydantic import validates_with_pydantic

def f(x: PositiveInt):
    return x

Conf = builds(f, zen_wrappers=validates_with_pydantic)

instantiate(Conf, x=10)  # ok
instantiate(Conf, x=-10)  # pydantic raises validation error

More Advanced Details

Wrappers are configurable

The wrappers themselves can be made configurable:

ConfWrapper = builds(some_wrapper, ....)

Conf = builds(f, zen_wrapper=ConfWrapper)

None can be used as a placeholder for wrappers

None can be be interspersed among wrappers; it will be ignored. This means that None can be used as a placeholder for wrappers that are going to be toggled on/off from the CLI, via config groups. See the next subsection for an example.

Wrappers can be specified via interpolated fields

One can specify interpolation strings as wrappers. This, in conjunction with zen_meta means that wrappers can be controlled via groups and from the CLI:

>>> Conf = builds(
...     faren_to_celsius,
...     zen_wrappers=["${..verbose}", "${..convert}"],
...     zen_meta=dict(verbose=None, convert=None),
... )
>>> instantiate(Conf, temp_f=32)
0.0

Let's "turn on" our two wrappers (this can be done from CLI with groups, but we'll do it manually)

def yell(func):
    print(f"HI {func}")
    return func

>>> Conf.verbose = just(yell)
>>> Conf.convert = just(change_celcius_to_kelvin)

>>> print(to_yaml(Conf))  # it is still serializable!
_target_: hydra_zen.funcs.zen_processing
_zen_target: __main__.faren_to_celsius
_zen_exclude:
- verbose
- convert
_zen_wrappers:
- ${..verbose}
- ${..convert}
verbose:
  _target_: hydra_zen.funcs.get_obj
  path: __main__.yell
convert:
  _target_: hydra_zen.funcs.get_obj
  path: __main__.change_celcius_to_kelvin

>>> instantiate(Conf, temp_f=32)
HI <function faren_to_celsius at 0x000002B9D4B81AF0>
273.15

Bells and Whistles

yaml Representations

zen_processing handles the process of importing wrapper-functions for us, so we don't need to use Just everywhere. This leads to tidy yamls

>>> def pp(cfg): print(to_yaml(cfg))
    
>>> pp(builds(dict, zen_wrappers=change_celcius_to_kelvin))
_target_: hydra_zen.funcs.zen_processing
_zen_target: builtins.dict
_zen_wrappers: __main__.change_celcius_to_kelvin

>>> pp(builds(dict, zen_wrappers=[change_celcius_to_kelvin, yell]))
_target_: hydra_zen.funcs.zen_processing
_zen_target: builtins.dict
_zen_wrappers:
- __main__.change_celcius_to_kelvin
- __main__.yell

We will do some nice auto-flattening of wrappers too:

# a list of one wrapper will get flattened to just the wrapper
>>> pp(builds(dict, zen_wrappers=[change_celcius_to_kelvin]))
_target_: hydra_zen.funcs.zen_processing
_zen_target: builtins.dict
_zen_wrappers: __main__.change_celcius_to_kelvin

# just-configs get flattened too
>>> pp(builds(dict, zen_wrappers=[just(change_celcius_to_kelvin), just(yell)]))
_target_: hydra_zen.funcs.zen_processing
_zen_target: builtins.dict
_zen_wrappers:
- __main__.change_celcius_to_kelvin
- __main__.yell

User-Friendly Warning About Mismatched Interpolation Levels

Using interpolated strings with zen_meta and zen_wrappers together is a natural fit. However, it can be tricky to figure out the right level of interpolation when you have things nested in a list. This can lead to pretty hard-to-find runtime errors.

Fortunately, we can easily detect the level of relative interpolation, and warn users – recommending the correct interpolation.

E.g.

>>> builds(dict, zen_wrappers="${...s}", zen_meta=dict(s=None))
UserWarning: Building: dict ..
A zen-wrapper is specified via the interpolated field, ${...s}, along with the meta-field name s, 
however it appears to point to the wrong level. It is likely you should change ${...s} to ${.s}
           
>>> builds(dict, zen_wrappers=["${.s}", "${.s}"], zen_meta=dict(s=None))
UserWarning: Building: dict ..
A zen-wrapper is specified via the interpolated field, ${.s}, along with the meta-field names, 
however it appears to point to the wrong level. It is likely you should change ${.s} to ${..s}                       

We make no attempt to predict things about absolute interpolations, only relative.

User-Friendly Annotations

Because zen_wrappers is a somewhat-sophisticated feature, it is important that users are given as much guidance as possible so that they catch bugs before they actually run their Hydra job. In addition to aggressive input validation, we provide incisive annotations too:

[jgbos]

That was still a big change, luckily was looking at this before you finished.

[rsokl]

That was still a big change, luckily was looking at this before you finished.

Thanks for reviewing all of it!

[jgbos]

Yeah, this is really stuff. But now I have deprecation warnings everywhere :)