Make literal and union work with raw from_dict csp, add option to go via pydantic

Signed-off-by: Nijat Khanbabayev <[email protected]>

Author: NeejWeej

csp/impl/struct.py

@@ -35,7 +35,11 @@ def __new__(cls, name, bases, dct):
                 # Lists need to be normalized too as potentially we need to add a boolean flag to use FastList
                 if v == FastList:
                     raise TypeError(f"{v} annotation is not supported without args")
-                if CspTypingUtils.is_generic_container(v) or CspTypingUtils.is_union_type(v):
+                if (
+                    CspTypingUtils.is_generic_container(v)
+                    or CspTypingUtils.is_union_type(v)
+                    or CspTypingUtils.is_literal_type(v)
+                ):
                     actual_type = ContainerTypeNormalizer.normalized_type_to_actual_python_type(v)
                     if CspTypingUtils.is_generic_container(actual_type):
                         raise TypeError(f"{v} annotation is not supported as a struct field [{actual_type}]")
@@ -147,6 +151,17 @@ def serializer(val, handler):
 
 
 class Struct(_csptypesimpl.PyStruct, metaclass=StructMeta):
+    @classmethod
+    def type_adapter(cls):
+        # Late import to avoid autogen issues
+        from pydantic import TypeAdapter
+
+        internal_type_adapter = getattr(cls, "_pydantic_type_adapter", None)
+        if internal_type_adapter:
+            return internal_type_adapter
+        cls._pydantic_type_adapter = TypeAdapter(cls)
+        return cls._pydantic_type_adapter
+
     @classmethod
     def metadata(cls, typed=False):
         if typed:
@@ -191,7 +206,8 @@ def _obj_from_python(cls, json, obj_type):
         if CspTypingUtils.is_generic_container(obj_type):
             if CspTypingUtils.get_origin(obj_type) in (typing.List, typing.Set, typing.Tuple, FastList):
                 return_type = ContainerTypeNormalizer.normalized_type_to_actual_python_type(obj_type)
-                (expected_item_type,) = obj_type.__args__
+                # We only take the first item, so like for a Tuple, we would ignore arguments after
+                expected_item_type = obj_type.__args__[0]
                 return_type = list if isinstance(return_type, list) else return_type
                 return return_type(cls._obj_from_python(v, expected_item_type) for v in json)
             elif CspTypingUtils.get_origin(obj_type) is typing.Dict:
@@ -206,6 +222,13 @@ def _obj_from_python(cls, json, obj_type):
                 return json
             else:
                 raise NotImplementedError(f"Can not deserialize {obj_type} from json")
+        elif CspTypingUtils.is_union_type(obj_type):
+            return json  ## no checks, just let it through
+        elif CspTypingUtils.is_literal_type(obj_type):
+            return_type = ContainerTypeNormalizer.normalized_type_to_actual_python_type(obj_type)
+            if isinstance(json, return_type):
+                return json
+            raise ValueError(f"Expected type {return_type} received {json.__class__}")
         elif issubclass(obj_type, Struct):
             if not isinstance(json, dict):
                 raise TypeError("Representation of struct as json is expected to be of dict type")
@@ -223,7 +246,9 @@ def _obj_from_python(cls, json, obj_type):
                 return obj_type(json)
 
     @classmethod
-    def from_dict(cls, json: dict):
+    def from_dict(cls, json: dict, use_pydantic: bool = False):
+        if use_pydantic:
+            return cls.type_adapter().validate_python(json)
         return cls._obj_from_python(json, cls)
 
     def to_dict_depr(self):

csp/impl/types/container_type_normalizer.py

@@ -81,21 +81,21 @@ def normalized_type_to_actual_python_type(cls, typ, level=0):
                 return [cls.normalized_type_to_actual_python_type(typ.__args__[0], level + 1), True]
             if origin is typing.List and level == 0:
                 return [cls.normalized_type_to_actual_python_type(typ.__args__[0], level + 1)]
-            if origin is typing.Literal:
-                # Import here to prevent circular import
-                from csp.impl.types.instantiation_type_resolver import UpcastRegistry
-
-                args = typing.get_args(typ)
-                typ = type(args[0])
-                for arg in args[1:]:
-                    typ = UpcastRegistry.instance().resolve_type(typ, type(arg), raise_on_error=False)
-                if typ:
-                    return typ
-                else:
-                    return object
             return cls._NORMALIZED_TYPE_MAPPING.get(CspTypingUtils.get_origin(typ), typ)
         elif CspTypingUtils.is_union_type(typ):
             return object
+        elif CspTypingUtils.is_literal_type(typ):
+            # Import here to prevent circular import
+            from csp.impl.types.instantiation_type_resolver import UpcastRegistry
+
+            args = typing.get_args(typ)
+            typ = type(args[0])
+            for arg in args[1:]:
+                typ = UpcastRegistry.instance().resolve_type(typ, type(arg), raise_on_error=False)
+            if typ:
+                return typ
+            else:
+                return object
         else:
             return typ
 

csp/impl/types/typing_utils.py

@@ -15,6 +15,25 @@ class FastList(typing.List, typing.Generic[T]):  # Need to inherit from Generic[
     def __init__(self):
         raise NotImplementedError("Can not init FastList class")
 
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type, handler):
+        from pydantic_core import core_schema
+
+        # Late import to not interfere with autogen
+        args = typing.get_args(source_type)
+        if args:
+            inner_type = args[0]
+            list_schema = handler.generate_schema(typing.List[inner_type])
+        else:
+            list_schema = handler.generate_schema(typing.List)
+
+        def create_instance(raw_data, validator):
+            if isinstance(raw_data, FastList):
+                return raw_data
+            return validator(raw_data)  # just return a list
+
+        return core_schema.no_info_wrap_validator_function(function=create_instance, schema=list_schema)
+
 
 class CspTypingUtils39:
     _ORIGIN_COMPAT_MAP = {list: typing.List, set: typing.Set, dict: typing.Dict, tuple: typing.Tuple}
@@ -23,7 +42,7 @@ class CspTypingUtils39:
 
     @classmethod
     def is_generic_container(cls, typ):
-        return isinstance(typ, cls._GENERIC_ALIASES) and typ.__origin__ is not typing.Union
+        return isinstance(typ, cls._GENERIC_ALIASES) and typ.__origin__ not in (typing.Union, typing.Literal)
 
     @classmethod
     def is_type_spec(cls, val):
@@ -56,6 +75,10 @@ def is_numpy_nd_array_type(cls, typ):
     def is_union_type(cls, typ):
         return isinstance(typ, typing._GenericAlias) and typ.__origin__ is typing.Union
 
+    @classmethod
+    def is_literal_type(cls, typ):
+        return isinstance(typ, typing._GenericAlias) and typ.__origin__ is typing.Literal
+
     @classmethod
     def is_forward_ref(cls, typ):
         return isinstance(typ, typing.ForwardRef)