Add support for TypeAliasType (#1214)

* Add support for TypeAliasType (new way to define type aliases in Python 3.12) * Fix nested type aliases * Add tests * fallback handling for older python versions. --------- Co-authored-by: Igor Sorokin <[email protected]> Co-authored-by: T. Franzel <[email protected]>
tfranzel · Mar 30, 2024 · 580bee1 · 580bee1
1 parent 06d3b47
commit 580bee1
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/drf_spectacular/plumbing.py b/drf_spectacular/plumbing.py
@@ -75,13 +75,18 @@ class Choices:  # type: ignore
 
 LITERAL_TYPES: Tuple[Any, ...] = ()
 TYPED_DICT_META_TYPES: Tuple[Any, ...] = ()
+TYPE_ALIAS_TYPES: Tuple[Any, ...] = ()
 
 if sys.version_info >= (3, 8):
     from typing import Literal as _PyLiteral
     from typing import _TypedDictMeta as _PyTypedDictMeta  # type: ignore[attr-defined]
     LITERAL_TYPES += (_PyLiteral,)
     TYPED_DICT_META_TYPES += (_PyTypedDictMeta,)
 
+if sys.version_info >= (3, 12):
+    from typing import TypeAliasType
+    TYPE_ALIAS_TYPES += (TypeAliasType,)
+
 try:
     from typing_extensions import Literal as _PxLiteral
     from typing_extensions import _TypedDictMeta as _PxTypedDictMeta  # type: ignore[attr-defined]
@@ -1347,6 +1352,8 @@ def resolve_type_hint(hint):
         return schema
     elif isinstance(hint, TYPED_DICT_META_TYPES):
         return _resolve_typeddict(hint)
+    elif isinstance(hint, TYPE_ALIAS_TYPES):
+        return resolve_type_hint(hint.__value__)
     elif origin in UNION_TYPES:
         type_args = [arg for arg in args if arg is not type(None)]  # noqa: E721
         if len(type_args) > 1:

diff --git a/tests/test_plumbing.py b/tests/test_plumbing.py
@@ -321,6 +321,26 @@ class TD4(TD4Optional):
         )
     ])
 
+if sys.version_info >= (3, 12):
+    exec("type MyAlias = typing.Literal['x', 'y']")
+    exec("type MyAliasNested = MyAlias | list[int | str]")
+
+    TYPE_HINT_TEST_PARAMS.extend([
+        (
+            MyAlias,  # noqa: F821
+            {'enum': ['x', 'y'], 'type': 'string'}
+        ),
+        (
+            MyAliasNested,  # noqa: F821
+            {
+                'oneOf': [
+                    {'enum': ['x', 'y'], 'type': 'string'},
+                    {"type": "array", "items": {"oneOf": [{"type": "integer"}, {"type": "string"}]}}
+                ]
+            }
+        )
+    ])
+
 
 @pytest.mark.parametrize(['type_hint', 'ref_schema'], TYPE_HINT_TEST_PARAMS)
 def test_type_hint_extraction(no_warnings, type_hint, ref_schema):