Support default values in typing.List[dataclass] and typing.Dict[str, dataclass]

flytekit/core/type_engine.py

@@ -360,6 +360,7 @@ def assert_type(self, expected_type: Type[DataClassJsonMixin], v: T):
 
         expected_type = get_underlying_type(expected_type)
         expected_fields_dict = {}
+
         for f in dataclasses.fields(expected_type):
             expected_fields_dict[f.name] = f.type
 
@@ -539,11 +540,13 @@ def _get_origin_type_in_annotation(self, python_type: Type[T]) -> Type[T]:
                 field.type = self._get_origin_type_in_annotation(field.type)
         return python_type
 
-    def _fix_structured_dataset_type(self, python_type: Type[T], python_val: typing.Any) -> T:
+    def _fix_structured_dataset_type(self, python_type: Type[T], python_val: typing.Any) -> T | None:
         # In python 3.7, 3.8, DataclassJson will deserialize Annotated[StructuredDataset, kwtypes(..)] to a dict,
         # so here we convert it back to the Structured Dataset.
         from flytekit.types.structured import StructuredDataset
 
+        if python_val is None:
+            return python_val
         if python_type == StructuredDataset and type(python_val) == dict:
             return StructuredDataset(**python_val)
         elif get_origin(python_type) is list:
@@ -575,9 +578,13 @@ def _make_dataclass_serializable(self, python_val: T, python_type: Type[T]) -> t
             return self._make_dataclass_serializable(python_val, get_args(python_type)[0])
 
         if hasattr(python_type, "__origin__") and get_origin(python_type) is list:
+            if python_val is None:
+                return None
             return [self._make_dataclass_serializable(v, get_args(python_type)[0]) for v in cast(list, python_val)]
 
         if hasattr(python_type, "__origin__") and get_origin(python_type) is dict:
+            if python_val is None:
+                return None
             return {
                 k: self._make_dataclass_serializable(v, get_args(python_type)[1])
                 for k, v in cast(dict, python_val).items()

flytekit/interaction/click_types.py

@@ -1,11 +1,12 @@
+import dataclasses
 import datetime
 import enum
 import json
 import logging
 import os
 import pathlib
 import typing
-from typing import cast
+from typing import cast, get_args
 
 import rich_click as click
 import yaml
@@ -22,6 +23,7 @@
 from flytekit.types.file import FlyteFile
 from flytekit.types.iterator.json_iterator import JSONIteratorTransformer
 from flytekit.types.pickle.pickle import FlytePickleTransformer
+from flytekit.types.schema.types import FlyteSchema
 
 
 def is_pydantic_basemodel(python_type: typing.Type) -> bool:
@@ -305,11 +307,48 @@ def convert(
         if value is None:
             raise click.BadParameter("None value cannot be converted to a Json type.")
 
+        FLYTE_TYPES = [FlyteFile, FlyteDirectory, StructuredDataset, FlyteSchema]
+
+        def has_nested_dataclass(t: typing.Type) -> bool:
+            """
+            Recursively checks whether the given type or its nested types contain any dataclass.
+
+            This function is typically called with a dictionary or list type and will return True if
+            any of the nested types within the dictionary or list is a dataclass.
+
+            Note:
+            - A single dataclass will return True.
+            - The function specifically excludes certain Flyte types like FlyteFile, FlyteDirectory,
+            StructuredDataset, and FlyteSchema from being considered as dataclasses. This is because
+            these types are handled separately by Flyte and do not need to be converted to dataclasses.
+
+            Args:
+                t (typing.Type): The type to check for nested dataclasses.
+
+            Returns:
+                bool: True if the type or its nested types contain a dataclass, False otherwise.
+            """
+
+            if dataclasses.is_dataclass(t):
+                # FlyteTypes is not supported now, we can support it in the future.
+                return t not in FLYTE_TYPES
+
+            return any(has_nested_dataclass(arg) for arg in get_args(t))
+
         parsed_value = self._parse(value, param)
 
         # We compare the origin type because the json parsed value for list or dict is always a list or dict without
         # the covariant type information.
         if type(parsed_value) == typing.get_origin(self._python_type) or type(parsed_value) == self._python_type:
+            if get_args(self._python_type) == ():
+                return parsed_value
+            elif isinstance(parsed_value, list) and has_nested_dataclass(get_args(self._python_type)[0]):
+                j = JsonParamType(get_args(self._python_type)[0])
+                return [j.convert(v, param, ctx) for v in parsed_value]
+            elif isinstance(parsed_value, dict) and has_nested_dataclass(get_args(self._python_type)[1]):
+                j = JsonParamType(get_args(self._python_type)[1])
+                return {k: j.convert(v, param, ctx) for k, v in parsed_value.items()}
+
             return parsed_value
 
         if is_pydantic_basemodel(self._python_type):

tests/flytekit/unit/cli/pyflyte/my_wf_input.json

@@ -42,6 +42,9 @@
     },
     "p": "None",
     "q": "tests/flytekit/unit/cli/pyflyte/testdata",
+    "r": [{"i": 1, "a": ["h", "e"]}],
+    "s": {"x": {"i": 1, "a": ["h", "e"]}},
+    "t": {"i": [{"i":1,"a":["h","e"]}]},
     "remote": "tests/flytekit/unit/cli/pyflyte/testdata",
     "image": "tests/flytekit/unit/cli/pyflyte/testdata"
 }

tests/flytekit/unit/cli/pyflyte/my_wf_input.yaml

@@ -30,5 +30,22 @@ o:
     - tests/flytekit/unit/cli/pyflyte/testdata/df.parquet
 p: 'None'
 q: tests/flytekit/unit/cli/pyflyte/testdata
+r:
+  - i: 1
+    a:
+      - h
+      - e
+s:
+  x:
+    i: 1
+    a:
+      - h
+      - e
+t:
+  i:
+    - i: 1
+      a:
+        - h
+        - e
 remote: tests/flytekit/unit/cli/pyflyte/testdata
 image: tests/flytekit/unit/cli/pyflyte/testdata

tests/flytekit/unit/cli/pyflyte/test_run.py

@@ -201,6 +201,12 @@ def test_pyflyte_run_cli(workflow_file):
             "Any",
             "--q",
             DIR_NAME,
+            "--r",
+            json.dumps([{"i": 1, "a": ["h", "e"]}]),
+            "--s",
+            json.dumps({"x": {"i": 1, "a": ["h", "e"]}}),
+            "--t",
+            json.dumps({"i": [{"i":1,"a":["h","e"]}]}),
         ],
         catch_exceptions=False,
     )

tests/flytekit/unit/cli/pyflyte/workflow.py

@@ -35,6 +35,9 @@ class MyDataclass(DataClassJsonMixin):
     i: int
     a: typing.List[str]
 
+@dataclass
+class NestedDataclass(DataClassJsonMixin):
+    i: typing.List[MyDataclass]
 
 class Color(enum.Enum):
     RED = "RED"
@@ -61,8 +64,11 @@ def print_all(
     o: typing.Dict[str, typing.List[FlyteFile]],
     p: typing.Any,
     q: FlyteDirectory,
+    r: typing.List[MyDataclass],
+    s: typing.Dict[str, MyDataclass],
+    t: NestedDataclass,
 ):
-    print(f"{a}, {b}, {c}, {d}, {e}, {f}, {g}, {h}, {i}, {j}, {k}, {l}, {m}, {n}, {o}, {p}, {q}")
+    print(f"{a}, {b}, {c}, {d}, {e}, {f}, {g}, {h}, {i}, {j}, {k}, {l}, {m}, {n}, {o}, {p}, {q}, {r}, {s}, {t}")
 
 
 @task
@@ -93,14 +99,17 @@ def my_wf(
     o: typing.Dict[str, typing.List[FlyteFile]],
     p: typing.Any,
     q: FlyteDirectory,
+    r: typing.List[MyDataclass],
+    s: typing.Dict[str, MyDataclass],
+    t: NestedDataclass,
     remote: pd.DataFrame,
     image: StructuredDataset,
     m: dict = {"hello": "world"},
 ) -> Annotated[StructuredDataset, subset_cols]:
     x = get_subset_df(df=remote)  # noqa: shown for demonstration; users should use the same types between tasks
     show_sd(in_sd=x)
     show_sd(in_sd=image)
-    print_all(a=a, b=b, c=c, d=d, e=e, f=f, g=g, h=h, i=i, j=j, k=k, l=l, m=m, n=n, o=o, p=p, q=q)
+    print_all(a=a, b=b, c=c, d=d, e=e, f=f, g=g, h=h, i=i, j=j, k=k, l=l, m=m, n=n, o=o, p=p, q=q, r=r, s=s, t=t)
     return x