Extend annotation -> type map to include py310 native union types

optimade/models/types.py

@@ -1,3 +1,4 @@
+from types import UnionType
 from typing import Annotated, Optional, Union, get_args
 
 from pydantic import Field
@@ -23,7 +24,7 @@
 
 AnnotatedType = type(ChemicalSymbol)
 OptionalType = type(Optional[str])
-UnionType = type(Union[str, int])
+_UnionType = type(Union[str, int])
 NoneType = type(None)
 
 
@@ -39,7 +40,7 @@ def _get_origin_type(annotation: type) -> type:
     """
     # If the annotation is a Union, get the first non-None type (this includes
     # Optional[T])
-    if isinstance(annotation, (OptionalType, UnionType)):
+    if isinstance(annotation, (OptionalType, UnionType, _UnionType)):
         for arg in get_args(annotation):
             if arg not in (None, NoneType):
                 annotation = arg
@@ -50,7 +51,7 @@ def _get_origin_type(annotation: type) -> type:
         annotation = get_args(annotation)[0]
 
     # Recursively unpack annotation, if it is a Union, Optional, or Annotated type
-    while isinstance(annotation, (OptionalType, UnionType, AnnotatedType)):
+    while isinstance(annotation, (OptionalType, UnionType, _UnionType, AnnotatedType)):
         annotation = _get_origin_type(annotation)
 
     # Special case for Literal

tests/models/test_types.py

@@ -0,0 +1,13 @@
+from typing import Annotated, Optional
+
+
+def test_origin_type():
+    from optimade.models.types import _get_origin_type
+
+    assert _get_origin_type(int | None) is int
+    assert _get_origin_type(str | None) is str
+    assert _get_origin_type(Optional[int]) is int
+    assert _get_origin_type(Optional[str]) is str
+    assert _get_origin_type(Annotated[int, "test"]) is int
+    assert _get_origin_type(Annotated[str, "test"]) is str
+    assert _get_origin_type(int | str | None) is int