Fix annotation adjustment for pipe unions

Signed-off-by: Pascal Tomecek <pascal.tomecek@cubistsystematic.com>

Author: ptomecek

ccflow/base.py

@@ -6,11 +6,9 @@
 import logging
 import pathlib
 import platform
-import typing
 from types import GenericAlias, MappingProxyType
-from typing import Any, Callable, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, get_args, get_origin
+from typing import Any, Callable, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, get_args, get_origin
 
-import typing_extensions
 from omegaconf import DictConfig
 from packaging import version
 from pydantic import (
@@ -100,35 +98,18 @@ def get_registry_dependencies(self, types: Optional[Tuple["ModelType"]] = None)
 # NOTE: For this logic to be removed, require https://github.com/pydantic/pydantic-core/pull/1478
 from pydantic._internal._model_construction import ModelMetaclass  # noqa: E402
 
-# Required for py38 compatibility
-# In python 3.8, get_origin(List[float]) returns list, but you can't call list[float] to retrieve the annotation
-# Furthermore, Annotated is part of typing_Extensions and get_origin(Annotated[str, ...]) returns str rather than Annotated
-_IS_PY38 = version.parse(platform.python_version()) < version.parse("3.9")
-# For a more complete list, see https://github.com/alexmojaki/eval_type_backport/blob/main/eval_type_backport/eval_type_backport.py
-_PY38_ORIGIN_MAP = {
-    tuple: typing.Tuple,
-    list: typing.List,
-    dict: typing.Dict,
-    set: typing.Set,
-    frozenset: typing.FrozenSet,
-    collections.abc.Callable: typing.Callable,
-    collections.abc.Iterable: typing.Iterable,
-    collections.abc.Mapping: typing.Mapping,
-    collections.abc.MutableMapping: typing.MutableMapping,
-    collections.abc.Sequence: typing.Sequence,
-}
+_IS_PY39 = version.parse(platform.python_version()) < version.parse("3.10")
 
 
 def _adjust_annotations(annotation):
     origin = get_origin(annotation)
-    if _IS_PY38:
-        origin = _PY38_ORIGIN_MAP.get(origin, origin)
-        if isinstance(annotation, typing_extensions._AnnotatedAlias):
-            args = annotation.__metadata__
-        else:
-            args = get_args(annotation)
-    else:
-        args = get_args(annotation)
+    args = get_args(annotation)
+    if not _IS_PY39:
+        from types import UnionType
+
+        if origin is UnionType:
+            origin = Union
+
     if isinstance(annotation, GenericAlias) or (inspect.isclass(annotation) and issubclass(annotation, PydanticBaseModel)):
         return SerializeAsAny[annotation]
     elif origin and args:
@@ -139,10 +120,6 @@ def _adjust_annotations(annotation):
             return ClassVar[_adjust_annotations(args[0])]
         else:
             try:
-                if _IS_PY38 and isinstance(annotation, typing_extensions._AnnotatedAlias):
-                    if origin != annotation:
-                        origin = _adjust_annotations(origin)
-                    return typing_extensions.Annotated[(origin,) + tuple(_adjust_annotations(arg) for arg in args)]
                 return origin[tuple(_adjust_annotations(arg) for arg in args)]
             except TypeError:
                 raise TypeError(f"Could not adjust annotations for {origin}")

ccflow/tests/test_base_serialize.py

@@ -1,7 +1,9 @@
+import platform
 import unittest
 from typing import ClassVar, Dict, List, Optional, Type, Union
 
 import numpy as np
+from packaging import version
 from pydantic import BaseModel as PydanticBaseModel, ConfigDict, ValidationError
 
 from ccflow import BaseModel, NDArray
@@ -205,20 +207,27 @@ def test_serialize_as_any(self):
         from pydantic import SerializeAsAny
         from pydantic.types import constr
 
+        if version.parse(platform.python_version()) >= version.parse("3.10"):
+            pipe_union = A | int
+        else:
+            pipe_union = Union[A, int]
+
         class MyNestedModel(BaseModel):
             a1: A
             a2: Optional[Union[A, int]]
             a3: Dict[str, Optional[List[A]]]
             a4: ClassVar[A]
             a5: Type[A]
             a6: constr(min_length=1)
+            a7: pipe_union
 
         target = {
             "a1": SerializeAsAny[A],
             "a2": Optional[Union[SerializeAsAny[A], int]],
             "a4": ClassVar[SerializeAsAny[A]],
             "a5": Type[A],
             "a6": constr(min_length=1),  # Uses Annotation
+            "a7": Union[SerializeAsAny[A], int],
         }
         target["a3"] = dict[str, Optional[list[SerializeAsAny[A]]]]
         annotations = MyNestedModel.__annotations__
@@ -228,3 +237,4 @@ class MyNestedModel(BaseModel):
         self.assertEqual(str(annotations["a4"]), str(target["a4"]))
         self.assertEqual(str(annotations["a5"]), str(target["a5"]))
         self.assertEqual(str(annotations["a6"]), str(target["a6"]))
+        self.assertEqual(str(annotations["a7"]), str(target["a7"]))