Repository: mit-ll-responsible-ai/hydra-zen
Pull request #283 (2022-06-15 04:40:40)

Improved signature parsing: support for Annotated, Paramspec, NewType, TypeVarTuple
Prior to this, the presence of `ParamSpecArgs` or `ParamSpecKwargs` in a signature being parsed by `builds` would cause an unhashable type error.

This PR fixes this bug and adds support for extracting information from `Annotated` and `NewType`, and adds support for extrapolating from variadic tuples (e.g. `Tuple[int, *Ts]`) 

Test cases are also added for `Protocol` and `TypeVar`

## Before

```python
# `sanitized_type` is applied to types by `hydra_zen.builds` when parsing target signatures

>>> sanitized_type(Annotated[int, range(10)])
Any

>>> sanitized_type(NewType("H", str))
Any

>>> sanitized_type(Tuple[*Ts])
Tuple[Any]

>>> sanitized_type(P.args)  # error
>>> sanitized_type(P.kwargs)  # error
```

## After

```python
>>> sanitized_type(Annotated[int, range(10)])
int

>>> sanitized_type(NewType("H", str))
str

>>> sanitized_type(Tuple[*Ts])
Tuple[Any, ...]

>>> sanitized_type(P.args)
Any

>>> sanitized_type(P.kwargs)
Any
```


Patch:
diff --git a/setup.cfg b/setup.cfg
index b21818ed..22c95fac 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -44,7 +44,7 @@ commands = pytest \
 [testenv:min-deps]  # test against minimum dependency versions
 deps = hydra-core==1.1.0
        omegaconf==2.1.1
-       typing-extensions==4.0.1
+       typing-extensions==4.1.0
        {[testenv]deps}
 basepython = python3.7
 
diff --git a/setup.py b/setup.py
index e6d10aa9..9f231f15 100644
--- a/setup.py
+++ b/setup.py
@@ -27,7 +27,7 @@ KEYWORDS = (
 )
 INSTALL_REQUIRES = [
     "hydra-core >= 1.1.0",
-    "typing-extensions >= 4.0.1",
+    "typing-extensions >= 4.1.0",
 ]
 TESTS_REQUIRE = [
     "pytest >= 3.8",
diff --git a/src/hydra_zen/structured_configs/_utils.py b/src/hydra_zen/structured_configs/_utils.py
index bcf9da06..562cd0de 100644
--- a/src/hydra_zen/structured_configs/_utils.py
+++ b/src/hydra_zen/structured_configs/_utils.py
@@ -25,7 +25,22 @@ from typing import (
 )
 
 from omegaconf import II
-from typing_extensions import Final, TypeGuard
+from typing_extensions import (
+    Annotated,
+    Final,
+    ParamSpecArgs,
+    ParamSpecKwargs,
+    TypeGuard,
+    Unpack,
+)
+
+try:
+    from typing_extensions import _AnnotatedAlias
+except ImportError:  # pragma: no cover
+    # Python 3.6
+    class _AnnotatedAlias:
+        ...
+
 
 from hydra_zen._compatibility import (
     HYDRA_SUPPORTED_PRIMITIVE_TYPES,
@@ -303,6 +318,14 @@ def sanitized_type(
     >>> sanitized_type(Dict[str, frozenset])
     Dict[str, Any]
     """
+    if hasattr(type_, "__supertype__"):
+        # is NewType
+        return sanitized_type(
+            type_.__supertype__,
+            primitive_only=primitive_only,
+            wrap_optional=wrap_optional,
+            nested=nested,
+        )
 
     # Warning: mutating `type_` will mutate the signature being inspected
     # Even calling deepcopy(`type_`) silently fails to prevent this.
@@ -310,6 +333,28 @@ def sanitized_type(
     no_nested_container = not HYDRA_SUPPORTS_NESTED_CONTAINER_TYPES
 
     if origin is not None:
+
+        # Support for Annotated[x, y]
+        # Python 3.9+
+        # # type_: Annotated[x, y]; origin -> Annotated; args -> (x, y)
+        if origin is Annotated:  # pragma: no cover
+            return sanitized_type(
+                get_args(type_)[0],
+                primitive_only=primitive_only,
+                wrap_optional=wrap_optional,
+                nested=nested,
+            )
+
+        # Python 3.7-3.8
+        # type_: Annotated[x, y]; origin -> x
+        if isinstance(type_, _AnnotatedAlias):
+            return sanitized_type(
+                origin,
+                primitive_only=primitive_only,
+                wrap_optional=wrap_optional,
+                nested=nested,
+            )
+
         if primitive_only:  # pragma: no cover
             return Any
 
@@ -356,13 +401,20 @@ def sanitized_type(
                 return Any  # bare Tuple not supported by hydra
             args = cast(Tuple[type, ...], args)
             unique_args = set(args)
+
+            if any(get_origin(tp) is Unpack for tp in unique_args):
+                # E.g. Tuple[*Ts]
+                return Tuple[Any, ...]
+
             has_ellipses = Ellipsis in unique_args
 
+            # E.g. Tuple[int, int, int] or Tuple[int, ...]
             _unique_type = (
                 sanitized_type(args[0], primitive_only=no_nested_container, nested=True)
                 if len(unique_args) == 1 or (len(unique_args) == 2 and has_ellipses)
                 else Any
             )
+
             if has_ellipses:
                 return Tuple[_unique_type, ...]  # type: ignore
             else:
@@ -373,6 +425,10 @@ def sanitized_type(
     if HYDRA_SUPPORTS_PARTIAL and isinstance(type_, type) and issubclass(type_, Path):
         type_ = Path
 
+    if isinstance(type_, (ParamSpecArgs, ParamSpecKwargs)):
+        # these aren't hashable -- can't check for membership in set
+        return Any
+
     if (
         type_ is Any
         or type_ in HYDRA_SUPPORTED_PRIMITIVE_TYPES


Test patch:
diff --git a/tests/test_third_party/test_type_validators.py b/tests/test_third_party/test_type_validators.py
index 51528b92..381b2a5a 100644
--- a/tests/test_third_party/test_type_validators.py
+++ b/tests/test_third_party/test_type_validators.py
@@ -211,7 +211,6 @@ class UserIdentity(TypedDict):
         (Union[List[str], str], dict(a=1)),
         (MyNamedTuple, (1.0, 2.0)),
         (xf.add(UserIdentity, BEAR), dict(first="bruce", last="lee")),
-        (Annotated[int, "special"], "hello"),
     ],
 )
 @pytest.mark.parametrize("validator", all_validators)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 8e840e81..5c23b6c5 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -7,7 +7,19 @@ import string
 import sys
 from dataclasses import dataclass, field as dataclass_field
 from pathlib import Path, PosixPath, WindowsPath
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, TypeVar, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    NewType,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
 
 import hypothesis.strategies as st
 import pytest
@@ -19,12 +31,23 @@ from omegaconf.errors import (
     ConfigValueError,
     KeyValidationError,
 )
-from typing_extensions import Final, Literal
+from typing_extensions import (
+    Annotated,
+    Final,
+    Literal,
+    ParamSpec,
+    Protocol,
+    Self,
+    TypeAlias,
+    TypeVarTuple,
+    Unpack,
+)
 
 from hydra_zen import builds, instantiate, mutable_value
 from hydra_zen._compatibility import (
     HYDRA_SUPPORTS_BYTES,
     HYDRA_SUPPORTS_NESTED_CONTAINER_TYPES,
+    HYDRA_VERSION,
     HYDRA_SUPPORTS_Path,
     Version,
     _get_version,
@@ -39,6 +62,8 @@ from hydra_zen.typing import Builds
 from tests import everything_except
 
 T = TypeVar("T")
+P = ParamSpec("P")
+Ts = TypeVarTuple("Ts")
 
 current_module: str = sys.modules[__name__].__name__
 
@@ -108,7 +133,11 @@ class Color(enum.Enum):
     pass
 
 
-NoneType = type(None)
+class SomeProtocol(Protocol[T]):
+    ...
+
+
+NoneType: TypeAlias = None
 
 
 @pytest.mark.parametrize(
@@ -132,6 +161,62 @@ NoneType = type(None)
         (callable, Any),
         (frozenset, Any),
         (T, Any),
+        (List[T], List[Any]),
+        (Tuple[T, T], Tuple[Any, Any]),
+        (Callable[P, int], Any),
+        (P, Any),
+        (P.args, Any),  # type: ignore
+        (P.kwargs, Any),  # type: ignore
+        (Ts, Any),
+        (SomeProtocol[T], Any),
+        pytest.param(
+            Tuple[Unpack[Ts]],
+            Tuple[Any, ...],
+            marks=pytest.mark.skipif(
+                sys.version_info <= (3, 7), reason="Python 3.6 doesn't support Unpack"
+            ),
+        ),
+        pytest.param(
+            Tuple[Unpack[Ts], int],
+            Tuple[Any, ...],
+            marks=pytest.mark.skipif(
+                sys.version_info < (3, 7), reason="Python 3.6 doesn't support Unpack"
+            ),
+        ),
+        pytest.param(
+            Tuple[str, Unpack[Ts]],
+            Tuple[Any, ...],
+            marks=[
+                pytest.mark.skipif(
+                    sys.version_info < (3, 7),
+                    reason="Python 3.6 doesn't support Unpack",
+                ),
+                pytest.mark.xfail(
+                    HYDRA_VERSION < Version(1, 2, 0),
+                    reason="Hydra 1.1.2 doesn't parse tuples beyond first element.",
+                ),
+            ],
+        ),
+        pytest.param(
+            Tuple[str, Unpack[Ts], int],
+            Tuple[Any, ...],
+            marks=[
+                pytest.mark.skipif(
+                    sys.version_info < (3, 7),
+                    reason="Python 3.6 doesn't support Unpack",
+                ),
+                pytest.mark.xfail(
+                    HYDRA_VERSION < Version(1, 2, 0),
+                    reason="Hydra 1.1.2 doesn't parse tuples beyond first element.",
+                ),
+            ],
+        ),
+        (Annotated[int, int], int),
+        (Annotated[Tuple[str, str], int], Tuple[str, str]),
+        (Annotated[Builds, int], Any),
+        (NewType("I", int), int),
+        (NewType("S", Tuple[str, str]), Tuple[str, str]),
+        (Self, Any),  # type: ignore
         (Literal[1, 2], Any),  # unsupported generics
         (Type[int], Any),
         (Builds, Any),
@@ -222,6 +307,7 @@ def test_sanitized_type_expected_behavior(in_type, expected_type):
                 KeyValidationError,
                 ConfigIndexError,
                 ConfigValueError,
+                AttributeError,
             )
         ):
             Conf = OmegaConf.create(Bad)  # type: ignore
@@ -231,7 +317,8 @@ def test_sanitized_type_expected_behavior(in_type, expected_type):
     class Tmp:
         x: expected_type
 
-    OmegaConf.create(Tmp)  # shouldn't raise on `expected_type`
+    # shouldn't raise on `expected_type`
+    OmegaConf.create(Tmp)  # type: ignore
 
 
 def test_tuple_annotation_normalization():