Repository: mit-ll-responsible-ai/hydra-zen
Pull request #198 (2022-01-13 03:49:41)

Support partial'd object as configured value
Partially addresses #197 

Now users can directly specify a "partial'd" object -- i.e. an object produced by `functools.partial(...)` -- directly as a configured value in `make_config` and `builds`. 

Consider the following.

```python
from hydra_zen import make_config, instantiate
from functools import partial

def f(x, y): 
    return (x, y)
```
Before

```python
>>> make_config(v=partial(f, x=1))
HydraZenUnsupportedPrimitiveError: ...
```

After:

```python
>>> Conf = make_config(v=partial(f, x=1))
>>> instantiate(Conf)
>>> {'v': functools.partial(<function f at 0x0000015924B6A700>, x=1)}
```

### Implementation Details

Under the hood, `make_config` and `builds` will identify `partial`-type configured values and "unpack" them, so that we can use `builds(..., zen_partial=True)` to create a valid targeted config that describes the same object.

The unpacking is simple:

```python
def _unpack_partial(value: Partial[_T]) -> Type[PartialBuilds[Type[_T]]]:
    target = cast(Type[_T], value.func)
    return builds(target, *value.args, **value.keywords, zen_partial=True)
```

### Some Additional Changes
`hydra_zen.typing.Partial` is updated to be made a protocol, rather than a generic type. Also, its definition was modified to more closely match the implementation of `functools.partial`. We now have tests to check that runtime and static checks using the protocol are OK.


Patch:
diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst
index b5404112..b7f50b66 100644
--- a/docs/source/api_reference.rst
+++ b/docs/source/api_reference.rst
@@ -148,6 +148,7 @@ hydra-zen provides such support for values of the following types:
 - :py:class:`complex`
 - :py:class:`collections.Counter`
 - :py:class:`collections.deque`
+- :py:func:`functools.partial`  (*added in v0.5.0*)
 - :py:class:`pathlib.Path`
 - :py:class:`pathlib.PosixPath`
 - :py:class:`pathlib.WindowsPath`
diff --git a/src/hydra_zen/_compatibility.py b/src/hydra_zen/_compatibility.py
index 471440dc..4c1bfef4 100644
--- a/src/hydra_zen/_compatibility.py
+++ b/src/hydra_zen/_compatibility.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: MIT
 from collections import Counter, deque
 from enum import Enum
+from functools import partial
 from pathlib import Path, PosixPath, WindowsPath
 from typing import NamedTuple, Optional, Set
 
@@ -54,6 +55,7 @@ ZEN_SUPPORTED_PRIMITIVES: Set[type] = {
     set,
     frozenset,
     complex,
+    partial,
     Path,
     PosixPath,
     WindowsPath,
diff --git a/src/hydra_zen/structured_configs/_implementations.py b/src/hydra_zen/structured_configs/_implementations.py
index a34671a6..6ece8afe 100644
--- a/src/hydra_zen/structured_configs/_implementations.py
+++ b/src/hydra_zen/structured_configs/_implementations.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2022 Massachusetts Institute of Technology
 # SPDX-License-Identifier: MIT
+import functools
 import inspect
 import sys
 import warnings
@@ -53,7 +54,14 @@ from hydra_zen.errors import (
 )
 from hydra_zen.funcs import get_obj, partial, zen_processing
 from hydra_zen.structured_configs import _utils
-from hydra_zen.typing import Builds, Importable, Just, PartialBuilds, SupportedPrimitive
+from hydra_zen.typing import (
+    Builds,
+    Importable,
+    Just,
+    Partial,
+    PartialBuilds,
+    SupportedPrimitive,
+)
 from hydra_zen.typing._implementations import (
     DataClass,
     HasPartialTarget,
@@ -116,24 +124,6 @@ _VAR_KEYWORD: Final = inspect.Parameter.VAR_KEYWORD
 _builtin_function_or_method_type = type(len)
 
 
-def _cast_via_tuple(dest_type: Type[_T]) -> Callable[[_T], Type[Builds[Type[_T]]]]:
-    def converter(value):
-        return builds(dest_type, tuple(value))
-
-    return converter
-
-
-ZEN_VALUE_CONVERSION[set] = _cast_via_tuple(set)
-ZEN_VALUE_CONVERSION[frozenset] = _cast_via_tuple(frozenset)
-ZEN_VALUE_CONVERSION[deque] = _cast_via_tuple(deque)
-ZEN_VALUE_CONVERSION[bytes] = _cast_via_tuple(bytes)
-ZEN_VALUE_CONVERSION[bytearray] = _cast_via_tuple(bytearray)
-ZEN_VALUE_CONVERSION[range] = lambda value: builds(
-    range, value.start, value.stop, value.step
-)
-ZEN_VALUE_CONVERSION[Counter] = lambda counter: builds(Counter, dict(counter))
-
-
 def _get_target(x):
     return getattr(x, _TARGET_FIELD_NAME)
 
@@ -2456,3 +2446,28 @@ def _repack_zenfield(value: ZenField, name: str, bases: Tuple[_DataClass, ...]):
             _permit_default_factory=False,
         )
     return value
+
+
+# registering value-conversions that depend on `builds`
+def _cast_via_tuple(dest_type: Type[_T]) -> Callable[[_T], Type[Builds[Type[_T]]]]:
+    def converter(value):
+        return builds(dest_type, tuple(value))
+
+    return converter
+
+
+def _unpack_partial(value: Partial[_T]) -> Type[PartialBuilds[Type[_T]]]:
+    target = cast(Type[_T], value.func)
+    return builds(target, *value.args, **value.keywords, zen_partial=True)
+
+
+ZEN_VALUE_CONVERSION[set] = _cast_via_tuple(set)
+ZEN_VALUE_CONVERSION[frozenset] = _cast_via_tuple(frozenset)
+ZEN_VALUE_CONVERSION[deque] = _cast_via_tuple(deque)
+ZEN_VALUE_CONVERSION[bytes] = _cast_via_tuple(bytes)
+ZEN_VALUE_CONVERSION[bytearray] = _cast_via_tuple(bytearray)
+ZEN_VALUE_CONVERSION[range] = lambda value: builds(
+    range, value.start, value.stop, value.step
+)
+ZEN_VALUE_CONVERSION[Counter] = lambda counter: builds(Counter, dict(counter))
+ZEN_VALUE_CONVERSION[functools.partial] = _unpack_partial
diff --git a/src/hydra_zen/typing/_implementations.py b/src/hydra_zen/typing/_implementations.py
index 9193618a..5adaa3c7 100644
--- a/src/hydra_zen/typing/_implementations.py
+++ b/src/hydra_zen/typing/_implementations.py
@@ -14,6 +14,7 @@ from typing import (
     NewType,
     Sequence,
     Tuple,
+    Type,
     TypeVar,
     Union,
 )
@@ -36,19 +37,22 @@ class EmptyDict(TypedDict):
 
 
 _T = TypeVar("_T", covariant=True)
+T2 = TypeVar("T2")
+T3 = TypeVar("T3")
 
 
-class Partial(Generic[_T]):
-    func: Callable[..., _T]
+@runtime_checkable
+class Partial(Protocol[T2]):
+    func: Callable[..., T2]
     args: Tuple[Any, ...]
     keywords: Dict[str, Any]
 
-    def __init__(
-        self, func: Callable[..., _T], *args: Any, **kwargs: Any
-    ) -> None:  # pragma: no cover
+    def __new__(
+        cls: Type[T3], func: Callable[..., T2], *args: Any, **kwargs: Any
+    ) -> T3:  # pragma: no cover
         ...
 
-    def __call__(self, *args: Any, **kwargs: Any) -> _T:  # pragma: no cover
+    def __call__(self, *args: Any, **kwargs: Any) -> T2:  # pragma: no cover
         ...
 
 


Test patch:
diff --git a/tests/annotations/declarations.py b/tests/annotations/declarations.py
index 6ad972a6..b8c606eb 100644
--- a/tests/annotations/declarations.py
+++ b/tests/annotations/declarations.py
@@ -4,6 +4,7 @@
 from collections import Counter, deque
 from dataclasses import dataclass
 from enum import Enum
+from functools import partial
 from pathlib import Path
 from typing import Callable, List, Literal, Tuple, Type, TypeVar
 
@@ -19,7 +20,7 @@ from hydra_zen import (
     make_custom_builds_fn,
     mutable_value,
 )
-from hydra_zen.typing import Builds
+from hydra_zen.typing import Builds, Partial
 from hydra_zen.typing._implementations import HydraPartialBuilds
 
 T = TypeVar("T")
@@ -350,3 +351,9 @@ def make_hydra_partial(x: T) -> HydraPartialBuilds[Type[T]]:
 def check_HydraPartialBuilds():
     cfg = make_hydra_partial(int)
     a: Literal["Partial[int]"] = reveal_type(instantiate(cfg))
+
+
+def check_partial_protocol():
+    x: Partial[int]
+    x = partial(int)
+    # x = partial(str)  # should fail
diff --git a/tests/test_protocols.py b/tests/test_protocols.py
index 3b379973..55b1cbc2 100644
--- a/tests/test_protocols.py
+++ b/tests/test_protocols.py
@@ -14,7 +14,7 @@ from hydra_zen.structured_configs._implementations import (
     is_just,
     is_partial_builds,
 )
-from hydra_zen.typing import Builds, Just, PartialBuilds
+from hydra_zen.typing import Builds, Just, Partial, PartialBuilds
 
 
 @pytest.mark.parametrize(
@@ -139,3 +139,8 @@ def test_protocol_checkers(x, yes_builds, yes_just, yes_partial):
 
     if yes_builds or yes_just or yes_partial:
         instantiate(x)
+
+
+def test_partial_protocol():
+    assert isinstance(partial(int), Partial)
+    assert not isinstance(print, Partial)
diff --git a/tests/test_value_conversion.py b/tests/test_value_conversion.py
index 1ce77e21..9dc8decb 100644
--- a/tests/test_value_conversion.py
+++ b/tests/test_value_conversion.py
@@ -3,6 +3,7 @@
 import string
 from collections import Counter, deque
 from enum import Enum
+from functools import partial
 from pathlib import Path
 from typing import Dict, FrozenSet, List, Set, Union
 
@@ -13,6 +14,7 @@ from omegaconf import DictConfig, ListConfig, OmegaConf
 
 from hydra_zen import builds, instantiate, make_config, to_yaml
 from hydra_zen._compatibility import ZEN_SUPPORTED_PRIMITIVES
+from hydra_zen.errors import HydraZenValidationError
 from hydra_zen.structured_configs._value_conversion import ZEN_VALUE_CONVERSION
 
 
@@ -116,3 +118,53 @@ def test_Counter_roundtrip(x):
     out_counter = instantiate(OmegaConf.structured(to_yaml(BuildsConf)))["counter"]
 
     assert counter == out_counter
+
+
+def f1(*args, **kwargs):
+    return
+
+
+def f2(*args, **kwargs):
+    return
+
+
+@given(
+    target=st.sampled_from([f1, f2]),
+    args=st.lists(st.integers() | st.text(string.ascii_letters)).map(tuple),
+    kwargs=st.dictionaries(
+        keys=st.text("abcd", min_size=1),
+        values=st.integers() | st.text(string.ascii_letters),
+    ),
+    as_builds=st.booleans(),
+    via_yaml=st.booleans(),
+)
+def test_functools_partial_as_configured_value(
+    target, args, kwargs, as_builds: bool, via_yaml: bool
+):
+    partiald_obj = partial(target, *args, **kwargs)
+    Conf = (
+        make_config(field=partiald_obj)
+        if not as_builds
+        else builds(dict, field=partiald_obj)
+    )
+
+    if via_yaml:
+        Conf = OmegaConf.structured(to_yaml(Conf))
+
+    out = instantiate(Conf)["field"]
+
+    assert isinstance(out, partial)
+    assert out.func is target
+    assert out.args == args
+    assert out.keywords == kwargs
+
+
+def f3(z):
+    return
+
+
+def test_functools_partial_gets_validated():
+    make_config(x=partial(f3, z=2))  # OK
+
+    with pytest.raises(TypeError):
+        make_config(x=partial(f3, y=2))  # no param named `y`