convert union type with none to optional

Author: Rusteam

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,6 +1,6 @@
 from inspect import isclass
 from types import GenericAlias
-from typing import Dict, Union, get_args, get_origin, get_type_hints
+from typing import Dict, Union, Optional, get_args, get_origin, get_type_hints
 
 from typing_extensions import TypedDict
 
@@ -36,6 +36,12 @@ def _handle_generics(t_) -> GenericAlias:
             else:
                 new_type[arg_name] = arg_type
         if new_type:
+            # because TypedDict can't handle union types with None
+            # rewrite them as Optional[type]
+            for arg_name, arg_type in new_type.items():
+                type_args = get_args(arg_type)
+                if len(type_args) == 2 and type_args[1] is type(None):
+                    new_type[arg_name] = Optional[type_args[0]]
             return TypedDict(type_.__name__, new_type)
 
         return type_