Repository: mit-ll-responsible-ai/hydra-zen
Pull request #297 (2022-08-19 13:18:15)

Add support for dict,tuple,list for parity with OmegaConf 2.2.3
This PR improves hydra-zen's type sanitization behavior for empty sequences, and it adds parity with omegaconf==2.2.3's support for bare `dict`, `list`, and `tuple` annotations.

Before:

```python
>>> from hydra_zen.structured_configs._utils import sanitized_type

>>> sanitized_type(typing.Tuple)
typing.Any

>>> sanitized_type(list)
typing.Any
>>> sanitized_type(tuple)
typing.Any
>>> sanitized_type(dict)
typing.Any
```

After:
```python
# For omegaconf < 2.2.3
>>> sanitized_type(typing.Tuple)
typing.Any
>>> sanitized_type(list)
typing.List
>>> sanitized_type(tuple)
typing.Any
>>> sanitized_type(dict)
typing.Dict

# For 2.2.3 <= omegaconf
>>> sanitized_type(typing.Tuple)
typing.Tuple
>>> sanitized_type(list)
list
>>> sanitized_type(tuple)
tuple
>>> sanitized_type(dict)
dict
```


Patch:
diff --git a/docs/source/changes.rst b/docs/source/changes.rst
index da30486a..3e999dcb 100644
--- a/docs/source/changes.rst
+++ b/docs/source/changes.rst
@@ -31,6 +31,7 @@ Improvements
 - The :ref:`automatic type refinement <type-support>` performed by :func:`~hydra_zen.builds` now has enhanced support for ``typing.Annotated``, ``typing.NewType``, and ``typing.TypeVarTuple``. (See :pull:`283`)
 - :func:`~hydra_zen.builds` no longer has restrictions on inheritance patterns involving `PartialBuilds`-type configs. (See :pull:`290`)
 - Two new utility functions were added to the public API: :func:`~hydra_zen.is_partial_builds` and :func:`~hydra_zen.uses_zen_processing`
+- Improved :ref:`automatic type refinement <type-support>` for bare sequence types, and parity with support for `dict`, `list`, and `tuple` as type annotations in omegaconf 2.2.3+. (See :pull:`297`)
 
 Bug Fixes
 ---------
diff --git a/src/hydra_zen/structured_configs/_utils.py b/src/hydra_zen/structured_configs/_utils.py
index 92565927..99c5af74 100644
--- a/src/hydra_zen/structured_configs/_utils.py
+++ b/src/hydra_zen/structured_configs/_utils.py
@@ -24,7 +24,6 @@ from typing import (
     overload,
 )
 
-
 from omegaconf import II, DictConfig, ListConfig
 from typing_extensions import (
     Annotated,
@@ -47,7 +46,9 @@ from hydra_zen._compatibility import (
     HYDRA_SUPPORTED_PRIMITIVE_TYPES,
     HYDRA_SUPPORTS_NESTED_CONTAINER_TYPES,
     HYDRA_SUPPORTS_PARTIAL,
+    OMEGACONF_VERSION,
     PATCH_OMEGACONF_830,
+    Version,
 )
 from hydra_zen.errors import HydraZenValidationError
 from hydra_zen.typing._implementations import DataClass_, Field, InterpStr
@@ -287,6 +288,14 @@ def get_obj_path(obj: Any) -> str:
 
 
 NoneType = type(None)
+_supported_types = HYDRA_SUPPORTED_PRIMITIVE_TYPES | {
+    list,
+    dict,
+    tuple,
+    List,
+    Tuple,
+    Dict,
+}
 
 
 def sanitized_type(
@@ -328,9 +337,16 @@ def sanitized_type(
             nested=nested,
         )
 
+    if OMEGACONF_VERSION < Version(2, 2, 3):  # pragma: no cover
+        try:
+            type_ = {list: List, tuple: Tuple, dict: Dict}.get(type_, type_)
+        except TypeError:
+            pass
+
     # Warning: mutating `type_` will mutate the signature being inspected
     # Even calling deepcopy(`type_`) silently fails to prevent this.
     origin = get_origin(type_)
+
     no_nested_container = not HYDRA_SUPPORTS_NESTED_CONTAINER_TYPES
 
     if origin is not None:
@@ -379,17 +395,18 @@ def sanitized_type(
             return Union[optional_type, NoneType]  # type: ignore
 
         if origin is list or origin is List:
-            return List[sanitized_type(args[0], primitive_only=no_nested_container, nested=True)] if args else type_  # type: ignore
+            if args:
+                return List[sanitized_type(args[0], primitive_only=no_nested_container, nested=True)]  # type: ignore
+            return List
 
         if origin is dict or origin is Dict:
-            return (
-                Dict[
-                    sanitized_type(args[0], primitive_only=True, nested=True),  # type: ignore
-                    sanitized_type(args[1], primitive_only=no_nested_container, nested=True),  # type: ignore
-                ]
-                if args
-                else type_
-            )
+            if args:
+                KeyType = sanitized_type(args[0], primitive_only=True, nested=True)
+                ValueType = sanitized_type(
+                    args[1], primitive_only=no_nested_container, nested=True
+                )
+                return Dict[KeyType, ValueType]  # type: ignore
+            return Dict
 
         if (origin is tuple or origin is Tuple) and not nested:
             # hydra silently supports tuples of homogenous types
@@ -399,7 +416,8 @@ def sanitized_type(
             #
             # Otherwise we preserve the annotation as accurately as possible
             if not args:
-                return Any  # bare Tuple not supported by hydra
+                return Any if OMEGACONF_VERSION < (2, 2, 3) else Tuple
+
             args = cast(Tuple[type, ...], args)
             unique_args = set(args)
 
@@ -432,11 +450,14 @@ def sanitized_type(
 
     if (
         type_ is Any
-        or type_ in HYDRA_SUPPORTED_PRIMITIVE_TYPES
+        or type_ in _supported_types
         or is_dataclass(type_)
         or (isinstance(type_, type) and issubclass(type_, Enum))
     ):
 
+        if sys.version_info[:2] == (3, 6) and type_ is Dict:  # pragma: no cover
+            type_ = Dict[Any, Any]
+
         if wrap_optional and type_ is not Any:  # pragma: no cover
             # normally get_type_hints automatically resolves Optional[...]
             # when None is set as the default, but this has been flaky
@@ -447,7 +468,14 @@ def sanitized_type(
         return type_
 
     # Needed to cover python 3.6 where __origin__ doesn't normalize to type
-    if not primitive_only and type_ in {List, Tuple, Dict}:  # pragma: no cover
+    if not primitive_only and type_ in {
+        List,
+        Tuple,
+        Dict,
+        list,
+        dict,
+        tuple,
+    }:  # pragma: no cover
         if wrap_optional and type_ is not Any:
             type_ = Optional[type_]  # type: ignore
         return type_


Test patch:
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 5c23b6c5..c0c37862 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -48,6 +48,7 @@ from hydra_zen._compatibility import (
     HYDRA_SUPPORTS_BYTES,
     HYDRA_SUPPORTS_NESTED_CONTAINER_TYPES,
     HYDRA_VERSION,
+    OMEGACONF_VERSION,
     HYDRA_SUPPORTS_Path,
     Version,
     _get_version,
@@ -65,8 +66,6 @@ T = TypeVar("T")
 P = ParamSpec("P")
 Ts = TypeVarTuple("Ts")
 
-current_module: str = sys.modules[__name__].__name__
-
 
 def pass_through(*args):
     return args
@@ -133,12 +132,15 @@ class Color(enum.Enum):
     pass
 
 
-class SomeProtocol(Protocol[T]):
+class SomeProtocol(Protocol[T]):  # type: ignore
     ...
 
 
 NoneType: TypeAlias = None
 
+vList = List[Any] if sys.version_info < (3, 8) else List
+vDict = Dict[Any, Any] if sys.version_info < (3, 8) else Dict
+
 
 @pytest.mark.parametrize(
     "in_type, expected_type",
@@ -155,11 +157,14 @@ NoneType: TypeAlias = None
         (C, Any),  # unsupported primitives
         (type(None), Any),
         (set, Any),
-        (list, Any),
-        (tuple, Any),
-        (dict, Any),
+        (list, (vList if OMEGACONF_VERSION < Version(2, 2, 3) else list)),
+        (tuple, (Any if OMEGACONF_VERSION < Version(2, 2, 3) else tuple)),
+        (dict, (vDict if OMEGACONF_VERSION < Version(2, 2, 3) else dict)),
         (callable, Any),
         (frozenset, Any),
+        (vList, vList),
+        (Tuple, (Any if OMEGACONF_VERSION < Version(2, 2, 3) else Tuple)),
+        (Dict, vDict),
         (T, Any),
         (List[T], List[Any]),
         (Tuple[T, T], Tuple[Any, Any]),
diff --git a/tests/test_utils/test_sequence_coercion.py b/tests/test_utils/test_sequence_coercion.py
index 6e4356e2..3adafc3b 100644
--- a/tests/test_utils/test_sequence_coercion.py
+++ b/tests/test_utils/test_sequence_coercion.py
@@ -112,7 +112,7 @@ def test_fuzz_sequence_coercion(annotation):
 @example(None)
 @example(ListConfig([1, 2]))
 # generate anything that hypothesis can generate!
-@given(st.from_type(type).flatmap(st.from_type))  # type: ignore
+@given(st.from_type(type).flatmap(st.from_type))
 def test_only_list_values_get_cast(input_val):
     try:
         assume(input_val == input_val)  # reject inputs who don't satisfy identity