Fix the type of optional[int] in dataclass (#1135)

Signed-off-by: Kevin Su <[email protected]>

flytekit/core/type_engine.py

@@ -491,7 +491,7 @@ def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) ->
     def _fix_val_int(self, t: typing.Type, val: typing.Any) -> typing.Any:
         if val is None:
             return val
-        if t == int:
+        if t == int or t == typing.Optional[int]:
             return int(val)
 
         if isinstance(val, list):
@@ -1369,6 +1369,11 @@ def convert_json_schema_to_python_class(schema: dict, schema_name) -> Type[datac
 def _get_element_type(element_property: typing.Dict[str, str]) -> Type[T]:
     element_type = element_property["type"]
     element_format = element_property["format"] if "format" in element_property else None
+
+    if type(element_type) == list:
+        # Element type of Optional[int] is [integer, None]
+        return typing.Optional[_get_element_type({"type": element_type[0]})]
+
     if element_type == "string":
         return str
     elif element_type == "integer":

tests/flytekit/unit/core/test_type_engine.py

@@ -157,6 +157,8 @@ class Bar(object):
     @dataclass_json
     @dataclass()
     class Foo(object):
+        u: typing.Optional[int]
+        v: typing.Optional[int]
         w: int
         x: typing.List[int]
         y: typing.Dict[str, str]
@@ -173,6 +175,7 @@ class Foo(object):
     foo_class = convert_json_schema_to_python_class(schema["definitions"], "FooSchema")
 
     guessed_pv = transformer.to_python_value(ctx, lv, expected_python_type=typing.List[foo_class])
+    print("=====")
     pv = transformer.to_python_value(ctx, lv, expected_python_type=typing.List[Foo])
     assert isinstance(guessed_pv, list)
     assert guessed_pv[0].u == pv[0].u
@@ -184,9 +187,7 @@ class Foo(object):
     assert type(guessed_pv[0].u) == int
     assert guessed_pv[0].v is None
     assert type(guessed_pv[0].w) == int
-    assert type(guessed_pv[0].z.v) == int
     assert type(guessed_pv[0].z.x) == float
-    assert guessed_pv[0].z.v == pv[0].z.v
     assert guessed_pv[0].z.y == pv[0].z.y
     assert guessed_pv[0].z.z == pv[0].z.z
     assert pv[0] == dataclass_from_dict(Foo, asdict(guessed_pv[0]))
@@ -1221,6 +1222,7 @@ def test_pass_annotated_to_downstream_tasks():
     """
     Test to confirm that the loaded dataframe is not affected and can be used in @dynamic.
     """
+
     # pandas dataframe hash function
     def hash_pandas_dataframe(df: pd.DataFrame) -> str:
         return str(pd.util.hash_pandas_object(df))