Merge pull request #681 from mit-ll-responsible-ai/defaultdict

autoconfig support for defaultdict

docs/source/api_reference.rst

@@ -226,6 +226,7 @@ hydra-zen provides specialized auto-config support for values of the following t
 - :py:class:`bytearray`
 - :py:class:`complex`
 - :py:class:`collections.Counter`
+- :py:class:`collections.defaultdict`
 - :py:class:`collections.deque`
 - :py:class:`datetime.timedelta`
 - :py:func:`functools.partial` (note: not compatible with pickling)

src/hydra_zen/_compatibility.py

@@ -1,6 +1,6 @@
 # Copyright (c) 2024 Massachusetts Institute of Technology
 # SPDX-License-Identifier: MIT
-from collections import Counter, deque
+from collections import Counter, defaultdict, deque
 from datetime import timedelta
 from enum import Enum
 from functools import partial
@@ -70,6 +70,7 @@ def _get_version(ver_str: str) -> Version:
         Counter,
         range,
         timedelta,
+        defaultdict,
     }
 )
 

src/hydra_zen/funcs.py

@@ -101,3 +101,13 @@ def zen_processing(
     if _zen_partial is True:
         return _functools.partial(obj, *args, **kwargs)
     return obj(*args, **kwargs)
+
+
+def as_default_dict(
+    dict_: _tp.Dict[_tp.Any, _tp.Any], *, default_factory: _tp.Any
+) -> _tp.DefaultDict[_tp.Any, _tp.Any]:
+    from collections import defaultdict
+
+    obj = defaultdict(default_factory)
+    obj.update(dict_)
+    return obj

src/hydra_zen/structured_configs/_implementations.py

@@ -4,7 +4,7 @@
 import inspect
 import sys
 import warnings
-from collections import Counter, deque
+from collections import Counter, defaultdict, deque
 from collections.abc import Collection
 from dataclasses import (  # use this for runtime checks
     MISSING,
@@ -72,7 +72,7 @@
     HydraZenUnsupportedPrimitiveError,
     HydraZenValidationError,
 )
-from hydra_zen.funcs import get_obj
+from hydra_zen.funcs import as_default_dict, get_obj
 from hydra_zen.structured_configs import _utils
 from hydra_zen.structured_configs._type_guards import safe_getattr
 from hydra_zen.typing import (
@@ -3750,6 +3750,30 @@ def __post_init__(self, CBuildsFn: Type[BuildsFn[Any]]) -> None:
         del CBuildsFn
 
 
+@dataclass(unsafe_hash=True)
+class ConfigFromDefaultDict:
+    dict_: Dict[Any, Any]
+    default_factory: Any = field(init=False)
+    CBuildsFn: InitVar[Type[BuildsFn[Any]]]
+    _target_: str = BuildsFn._get_obj_path(as_default_dict)
+
+    def __post_init__(self, CBuildsFn: Type[BuildsFn[Any]]) -> None:
+        assert isinstance(self.dict_, defaultdict)
+        self.default_factory = CBuildsFn.just(self.dict_.default_factory)
+        out = CBuildsFn._make_hydra_compatible(
+            dict(self.dict_),
+            convert_dataclass=True,
+            allow_zen_conversion=True,
+            structured_conf_permitted=True,
+        )
+        assert isinstance(out, dict)
+        self.dict_ = out
+
+
+ZEN_VALUE_CONVERSION[defaultdict] = lambda dict_, CBuildsFn: ConfigFromDefaultDict(
+    dict_, CBuildsFn
+)
+
 ZEN_VALUE_CONVERSION[set] = partial(
     ConfigFromTuple, _target_=BuildsFn._get_obj_path(set)
 )

tests/test_value_conversion.py

@@ -3,7 +3,7 @@
 import inspect
 import pickle
 import string
-from collections import Counter, deque
+from collections import Counter, defaultdict, deque
 from dataclasses import dataclass, field
 from datetime import timedelta
 from enum import Enum
@@ -142,6 +142,9 @@ def test_value_supported_via_config_maker_functions(
 Brange = pik_blds(dict, x=range(1, 10, 3), zen_dataclass={"cls_name": "Brange"})
 Brange2 = pik_blds(dict, x=range(2), zen_dataclass={"cls_name": "Brange2"})
 Bcounter = pik_blds(dict, x=Counter("apple"), zen_dataclass={"cls_name": "Bcounter"})
+x = defaultdict(list)
+x.update({1: [1, 2]})
+Bdefaultdict = pik_blds(dict, x=x, zen_dataclass={"cls_name": "Bdefaultdict"})
 
 
 @pytest.mark.parametrize(
@@ -164,6 +167,7 @@ def test_value_supported_via_config_maker_functions(
         Brange,
         Brange2,
         Bcounter,
+        Bdefaultdict,
     ],
 )
 def test_pickle_compatibility(Config):
@@ -190,6 +194,7 @@ def test_pickle_compatibility(Config):
         Brange,
         Brange2,
         Bcounter,
+        Bdefaultdict,
     ],
 )
 def test_unsafe_hash_default(Config):
@@ -412,3 +417,13 @@ def test_known_failcase_hydra_2350():
     actual = instantiate(Conf, x={"b": 2})
     expected = {"b": 2}
     assert actual == expected, actual
+
+
+def test_default_dict():
+    x = defaultdict(list)
+    x.update({1: [1 + 2j, 2]})
+    Conf = builds(dict, x=x)
+    actual = instantiate(Conf)["x"]
+    assert actual == x
+    assert isinstance(actual, defaultdict)
+    assert actual.default_factory is list