Merge pull request #31 from Rusteam/fix-typing

convert union type with none to optional

src/hayhooks/server/pipelines/models.py

@@ -26,7 +26,11 @@ def get_request_model(pipeline_name: str, pipeline_inputs):
     for component_name, inputs in pipeline_inputs.items():
         component_model = {}
         for name, typedef in inputs.items():
-            input_type = handle_unsupported_types(typedef["type"], {DataFrame: dict})
+            try:
+                input_type = handle_unsupported_types(typedef["type"], {DataFrame: dict})
+            except TypeError as e:
+                print(f"ERROR at {component_name!r}, {name}: {typedef}")
+                raise e
             component_model[name] = (
                 input_type,
                 typedef.get("default_value", ...),
@@ -70,7 +74,10 @@ def convert_component_output(component_output):
     """
     result = {}
     for output_name, data in component_output.items():
-        get_value = lambda data: data.to_dict()["init_parameters"] if hasattr(data, "to_dict") else data
+
+        def get_value(data):
+            return data.to_dict()["init_parameters"] if hasattr(data, "to_dict") else data
+
         if type(data) is list:
             result[output_name] = [get_value(d) for d in data]
         else:

src/hayhooks/server/utils/create_valid_type.py

@@ -1,8 +1,6 @@
 from inspect import isclass
 from types import GenericAlias
-from typing import Dict, Union, get_args, get_origin, get_type_hints
-
-from typing_extensions import TypedDict
+from typing import Dict, Optional, Union, get_args, get_origin, get_type_hints
 
 
 def handle_unsupported_types(type_: type, types_mapping: Dict[type, type]) -> Union[GenericAlias, type]:
@@ -26,7 +24,13 @@ def _handle_generics(t_) -> GenericAlias:
             else:
                 result = t
             child_typing.append(result)
-        return GenericAlias(get_origin(t_), tuple(child_typing))
+
+        if len(child_typing) == 2 and child_typing[1] is type(None):
+            # because TypedDict can't handle union types with None
+            # rewrite them as Optional[type]
+            return Optional[child_typing[0]]
+        else:
+            return GenericAlias(get_origin(t_), tuple(child_typing))
 
     if isclass(type_):
         new_type = {}
@@ -35,8 +39,6 @@ def _handle_generics(t_) -> GenericAlias:
                 new_type[arg_name] = _handle_generics(arg_type)
             else:
                 new_type[arg_name] = arg_type
-        if new_type:
-            return TypedDict(type_.__name__, new_type)
 
         return type_