Repository: mit-ll-responsible-ai/hydra-zen
Pull request #259 (2022-04-11 14:31:25)

`to_yaml` and `save_yaml` now reflect hydra-zen's extended type support
Previously, creating/saving the yaml representation for [values of types supported specifically by hydra-zen](https://mit-ll-responsible-ai.github.io/hydra-zen/api_reference.html#additional-types-supported-via-hydra-zen) required a user to manually apply `just` to those values to convert them to structured configs:

Before:

```python
>>> from functools import partial
>>> from hydra_zen import just, to_yaml
>>> print(just(to_yaml({"a": [1+2j, partial(int, 3)]})))  # requires call to `just`
a:
- real: 1.0
  imag: 2.0
  _target_: builtins.complex
- _target_: builtins.int
  _partial_: true
  _args_:
  - 3
```

Now, `to_yaml` and `save_yaml` will automatically apply `just` when appropriate.

**After**

```python
>>> from hydra_zen import just, to_yaml
>>> print(to_yaml({"a": [1+2j, partial(int, 3)]}))  # no call to `just`
a:
- real: 1.0
  imag: 2.0
  _target_: builtins.complex
- _target_: builtins.int
  _partial_: true
  _args_:
  - 3
```

If [we enable users to register support for custom types](https://github.com/mit-ll-responsible-ai/hydra-zen/issues/257), then instances of those types will automatically be compatible with these utility functions

```python
from hydra_zen import register_converter, builds, just, save_yaml

class WeirdClass:
    def __init__(self, x: int):
        self.x = x + 2

    def __as_config__(self):
        # tells `hydra_zen.just` how to convert instances of 
        # `WeirdClass` to structured configs
        return builds(type(self), x=x - 2)
```

Now instances of `WeirdClass` can be directly saved to yaml files:

```python
inst = WeirdClass(x=100)

save_yaml(inst, "weird.yaml")  # just works!
```


Patch:
diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst
index 333d6303..b7603dcf 100644
--- a/docs/source/api_reference.rst
+++ b/docs/source/api_reference.rst
@@ -123,10 +123,8 @@ Values of the following types can be specified directly in configs:
 Additional Types, Supported via hydra-zen
 *****************************************
 
-    This feature was introduced in hydra-zen ``v0.4.0``
-
 Values of additional types can be specified directly via hydra-zen's 
-:ref:`config-creation functions <create-config>`, and hydra-zen will automatically 
+:ref:`config-creation functions <create-config>` (and other utility functions like ``to_yaml``), and hydra-zen will automatically 
 create targeted configs to represent those values in a way that is compatible 
 with Hydra. For example, a :py:class:`complex` value can be specified directly via :func:`~hydra_zen.make_config`, and a targeted config will be created for that value.
 
@@ -139,6 +137,13 @@ with Hydra. For example, a :py:class:`complex` value can be specified directly v
      real: 2.0
      imag: 3.0
      _target_: builtins.complex
+   
+   >>> from functools import partial
+   >>> print(to_yaml(partial(int, 3)))
+   _target_: builtins.int
+   _partial_: true
+   _args_:
+   - 3
 
 hydra-zen provides specialized support for values of the following types:
 
diff --git a/docs/source/changes.rst b/docs/source/changes.rst
index db08fd2b..b1d6c963 100644
--- a/docs/source/changes.rst
+++ b/docs/source/changes.rst
@@ -8,6 +8,33 @@ Changelog
 This is a record of all past hydra-zen releases and what went into them, in reverse 
 chronological order. All previous releases should still be available on pip.
 
+.. _v0.7.0:
+
+------------------
+0.7.0 - 2022-XX-XX
+------------------
+
+New Features
+------------
+:func:`~hydra_zen.just`, :func:`~hydra_zen.to_yaml`, and :func:`~hydra_zen.save_as_yaml` can directly 
+operate on values of :ref:`types with specialized support from hydra-zen <additional-types>`; these 
+values will automatically be converted to structured configs. 
+
+.. code-block:: pycon
+
+   >>> from functools import partial
+   >>> from hydra_zen import to_yamlm just
+
+   >>> def f(x: int): returns x**2
+   
+   >>> just(partial(f, x=2))
+   PartialBuilds_f(_target_='__main__.f', _partial_=True, x=2)
+
+   >>> print(to_yaml(partial(f, x=2)))
+   _target_: __main__.f
+   _partial_: true
+   x: 2
+
 .. _v0.6.0:
 
 ------------------
diff --git a/src/hydra_zen/_hydra_overloads.py b/src/hydra_zen/_hydra_overloads.py
index ea1b5a0e..308df6ef 100644
--- a/src/hydra_zen/_hydra_overloads.py
+++ b/src/hydra_zen/_hydra_overloads.py
@@ -21,11 +21,14 @@ E.g.
 # pyright: strict
 
 import pathlib
-from typing import IO, Any, Callable, Type, TypeVar, Union, overload
+from dataclasses import is_dataclass
+from functools import wraps
+from typing import IO, Any, Callable, Type, TypeVar, Union, cast, overload
 
 from hydra.utils import instantiate as hydra_instantiate
 from omegaconf import MISSING, DictConfig, ListConfig, OmegaConf
 
+from .structured_configs._just import just
 from .structured_configs._value_conversion import ConfigComplex, ConfigPath
 from .typing import Builds, Just, Partial
 from .typing._implementations import DataClass_, HasTarget, InstOrType, IsPartial
@@ -34,6 +37,7 @@ __all__ = ["instantiate", "to_yaml", "save_as_yaml", "load_from_yaml", "MISSING"
 
 
 T = TypeVar("T")
+F = TypeVar("F", bound=Callable[..., Any])
 
 
 @overload
@@ -196,6 +200,17 @@ def instantiate(config: Any, *args: Any, **kwargs: Any) -> Any:
     return hydra_instantiate(config, *args, **kwargs)
 
 
+def _apply_just(fn: F) -> F:
+    @wraps(fn)
+    def wrapper(cfg: Any, *args: Any, **kwargs: Any):
+        if not is_dataclass(cfg):
+            cfg = just(cfg)
+        return fn(cfg, *args, **kwargs)
+
+    return cast(F, wrapper)
+
+
+@_apply_just
 def to_yaml(cfg: Any, *, resolve: bool = False, sort_keys: bool = False) -> str:
     """
     Serialize a config as a yaml-formatted string.
@@ -205,7 +220,7 @@ def to_yaml(cfg: Any, *, resolve: bool = False, sort_keys: bool = False) -> str:
     Parameters
     ----------
     cfg : Any
-        A valid configuration object (e.g. DictConfig, ListConfig, dataclass).
+        A valid configuration object, supported either by Hydra or hydra-zen
 
     resolve : bool, optional (default=False)
         If `True`, interpolated fields in `cfg` will be resolved in the yaml.
@@ -243,6 +258,13 @@ def to_yaml(cfg: Any, *, resolve: bool = False, sort_keys: bool = False) -> str:
     _target_: builtins.dict
     'y': 10
 
+    hydra-zen's additional supported types can be specified as well
+
+    >>> print(to_yaml(1+2j))
+    real: 1.0
+    imag: 2.0
+    _target_: builtins.complex
+
     **Specifying resolve**
 
     The following is a config with interpolated fields.
@@ -274,9 +296,11 @@ def to_yaml(cfg: Any, *, resolve: bool = False, sort_keys: bool = False) -> str:
     a: ???
     b: ???
     """
+
     return OmegaConf.to_yaml(cfg=cfg, resolve=resolve, sort_keys=sort_keys)
 
 
+@_apply_just
 def save_as_yaml(
     config: Any, f: Union[str, pathlib.Path, IO[Any]], resolve: bool = False
 ) -> None:
diff --git a/src/hydra_zen/structured_configs/_just.py b/src/hydra_zen/structured_configs/_just.py
index a5306aad..f355ae68 100644
--- a/src/hydra_zen/structured_configs/_just.py
+++ b/src/hydra_zen/structured_configs/_just.py
@@ -98,8 +98,8 @@ def just(obj: Any) -> Any:
 
     >>> from hydra_zen import just, instantiate, to_yaml
 
-    Calling ``just`` on a type or function object will create a structured config
-    that, when instantiated, returns that object.
+    Using ``just`` to create a config that returns an object – without calling it –
+    upon instantiation.
 
     >>> class A: pass
     >>> Conf_A = just(A)  # returns a dataclass-type
@@ -108,7 +108,7 @@ def just(obj: Any) -> Any:
 
     >>> def my_func(x: int): pass
     >>> Conf_func = just(my_func)  # returns a dataclass-type
-    >>> instantiate(Conf_func) is my_funct
+    >>> instantiate(Conf_func) is my_func
     True
 
 


Test patch:
diff --git a/tests/test_hydra_overloads.py b/tests/test_hydra_overloads.py
index 67e93b50..77df830f 100644
--- a/tests/test_hydra_overloads.py
+++ b/tests/test_hydra_overloads.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: MIT
 
 import string
+from pathlib import Path
 
 import hypothesis.strategies as st
 import pytest
@@ -12,6 +13,7 @@ from hydra_zen import (
     MISSING as zen_MISSING,
     builds,
     instantiate,
+    just,
     load_from_yaml,
     save_as_yaml,
     to_yaml,
@@ -45,3 +47,15 @@ def test_yaml_io_roundtrip(fname, value):
     loaded_conf = load_from_yaml(fname)
 
     assert instantiate(conf) == instantiate(loaded_conf)
+
+
+def test_to_yaml_applies_just():
+    data = [1 + 2j, Path.cwd(), "hi"]
+    assert to_yaml(just(data)) == to_yaml(data)
+
+
+def test_save_yaml_applies_just(tmp_path):
+    data = {"a": [1 + 2j, "hi", 3j]}
+    save_as_yaml(data, tmp_path / "cfg.yml")
+    out = dict(instantiate(load_from_yaml(tmp_path / "cfg.yml")))
+    assert out == data
diff --git a/tests/test_just.py b/tests/test_just.py
index 6e2f1327..5f38a771 100644
--- a/tests/test_just.py
+++ b/tests/test_just.py
@@ -3,7 +3,7 @@ from pathlib import Path
 
 import pytest
 
-from hydra_zen import instantiate, just
+from hydra_zen import instantiate, just, make_config
 from tests import is_same
 
 
@@ -44,3 +44,8 @@ def test_just_roundtrip(obj):
         assert is_same(out, obj)
     else:
         assert out == obj
+
+
+def test_just_of_structured_config_is_identity():
+    cfg = make_config(x=1)
+    assert just(cfg) is cfg