[Core feature] Allow dictionaries to be passed to execute workflows that take dataclasses

flytekit/core/type_engine.py

@@ -6,6 +6,7 @@
 import datetime as _datetime
 import enum
 import inspect
+import json
 import json as _json
 import mimetypes
 import textwrap
@@ -349,20 +350,61 @@ def assert_type(self, expected_type: Type[DataClassJsonMixin], v: T):
         for f in dataclasses.fields(expected_type):
             expected_fields_dict[f.name] = f.type
 
-        for f in dataclasses.fields(type(v)):  # type: ignore
-            original_type = f.type
-            expected_type = expected_fields_dict[f.name]
+        if isinstance(v, dict):
+            original_dict = v
 
-            if UnionTransformer.is_optional_type(original_type):
-                original_type = UnionTransformer.get_sub_type_in_optional(original_type)
-            if UnionTransformer.is_optional_type(expected_type):
-                expected_type = UnionTransformer.get_sub_type_in_optional(expected_type)
+            # Find the Optional keys in expected_fields_dict
+            optional_keys = {k for k, t in expected_fields_dict.items() if UnionTransformer.is_optional_type(t)}
 
-            val = v.__getattribute__(f.name)
-            if dataclasses.is_dataclass(val):
-                self.assert_type(expected_type, val)
-            elif original_type != expected_type:
-                raise TypeTransformerFailedError(f"Type of Val '{original_type}' is not an instance of {expected_type}")
+            # Remove the Optional keys from the keys of original_dict
+            original_key = set(original_dict.keys()) - optional_keys
+            expected_key = set(expected_fields_dict.keys()) - optional_keys
+
+            # Check if original_key is missing any keys from expected_key
+            missing_keys = expected_key - original_key
+            if missing_keys:
+                raise TypeTransformerFailedError(
+                    f"The original fields are missing the following keys from the dataclass fields: {list(missing_keys)}"
+                )
+
+            # Check if original_key has any extra keys that are not in expected_key
+            extra_keys = original_key - expected_key
+            if extra_keys:
+                raise TypeTransformerFailedError(
+                    f"The original fields have the following extra keys that are not in dataclass fields: {list(extra_keys)}"
+                )
+
+            for k, v in original_dict.items():
+                if k in expected_fields_dict:
+                    if isinstance(v, dict):
+                        self.assert_type(expected_fields_dict[k], v)
+                    else:
+                        expected_type = expected_fields_dict[k]
+                        original_type = type(v)
+                        if UnionTransformer.is_optional_type(expected_type):
+                            expected_type = UnionTransformer.get_sub_type_in_optional(expected_type)
+                        if original_type != expected_type:
+                            raise TypeTransformerFailedError(
+                                f"Type of Val '{original_type}' is not an instance of {expected_type}"
+                            )
+
+        else:
+            for f in dataclasses.fields(type(v)):  # type: ignore
+                original_type = f.type
+                expected_type = expected_fields_dict[f.name]
+
+                if UnionTransformer.is_optional_type(original_type):
+                    original_type = UnionTransformer.get_sub_type_in_optional(original_type)
+                if UnionTransformer.is_optional_type(expected_type):
+                    expected_type = UnionTransformer.get_sub_type_in_optional(expected_type)
+
+                val = v.__getattribute__(f.name)
+                if dataclasses.is_dataclass(val):
+                    self.assert_type(expected_type, val)
+                elif original_type != expected_type:
+                    raise TypeTransformerFailedError(
+                        f"Type of Val '{original_type}' is not an instance of {expected_type}"
+                    )
 
     def get_literal_type(self, t: Type[T]) -> LiteralType:
         """
@@ -424,6 +466,10 @@ def get_literal_type(self, t: Type[T]) -> LiteralType:
         return _type_models.LiteralType(simple=_type_models.SimpleType.STRUCT, metadata=schema, structure=ts)
 
     def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal:
+        if isinstance(python_val, dict):
+            json_str = json.dumps(python_val)
+            return Literal(scalar=Scalar(generic=_json_format.Parse(json_str, _struct.Struct())))
+
         if not dataclasses.is_dataclass(python_val):
             raise TypeTransformerFailedError(
                 f"{type(python_val)} is not of type @dataclass, only Dataclasses are supported for "

tests/flytekit/unit/core/test_type_engine.py

@@ -2,6 +2,7 @@
 import datetime
 import json
 import os
+import re
 import sys
 import tempfile
 import typing
@@ -1412,6 +1413,82 @@ class Bar(DataClassJsonMixin):
         DataclassTransformer().assert_type(gt, pv)
 
 
+@pytest.mark.skipif("pandas" not in sys.modules, reason="Pandas is not installed.")
+def test_assert_dict_type():
+    import pandas as pd
+
+    @dataclass
+    class AnotherDataClass(DataClassJsonMixin):
+        z: int
+
+    @dataclass
+    class Args(DataClassJsonMixin):
+        x: int
+        y: typing.Optional[str]
+        file: FlyteFile
+        dataset: StructuredDataset
+        another_dataclass: AnotherDataClass
+
+    pv = tempfile.mkdtemp(prefix="flyte-")
+    df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]})
+    sd = StructuredDataset(dataframe=df, file_format="parquet")
+    # Test when v is a dict
+    vd = {"x": 3, "y": "hello", "file": FlyteFile(pv), "dataset": sd, "another_dataclass": {"z": 4}}
+    DataclassTransformer().assert_type(Args, vd)
+
+    # Test when v is a dict but missing Optional keys and other keys from dataclass
+    md = {"x": 3, "file": FlyteFile(pv), "dataset": sd, "another_dataclass": {"z": 4}}
+    DataclassTransformer().assert_type(Args, md)
+
+    # Test when v is a dict but missing non-Optional keys from dataclass
+    md = {"y": "hello", "file": FlyteFile(pv), "dataset": sd, "another_dataclass": {"z": 4}}
+    with pytest.raises(
+        TypeTransformerFailedError,
+        match=re.escape("The original fields are missing the following keys from the dataclass fields: ['x']"),
+    ):
+        DataclassTransformer().assert_type(Args, md)
+
+    # Test when v is a dict but has extra keys that are not in dataclass
+    ed = {"x": 3, "y": "hello", "file": FlyteFile(pv), "dataset": sd, "another_dataclass": {"z": 4}, "z": "extra"}
+    with pytest.raises(
+        TypeTransformerFailedError,
+        match=re.escape("The original fields have the following extra keys that are not in dataclass fields: ['z']"),
+    ):
+        DataclassTransformer().assert_type(Args, ed)
+
+    # Test when the type of value in the dict does not match the expected_type in the dataclass
+    td = {"x": "3", "y": "hello", "file": FlyteFile(pv), "dataset": sd, "another_dataclass": {"z": 4}}
+    with pytest.raises(
+        TypeTransformerFailedError, match="Type of Val '<class 'str'>' is not an instance of <class 'int'>"
+    ):
+        DataclassTransformer().assert_type(Args, td)
+
+
+def test_to_literal_dict():
+    @dataclass
+    class Args(DataClassJsonMixin):
+        x: int
+        y: typing.Optional[str]
+
+    ctx = FlyteContext.current_context()
+    python_type = Args
+    expected = TypeEngine.to_literal_type(python_type)
+
+    # Test when python_val is a dict
+    python_val = {"x": 3, "y": "hello"}
+    literal = DataclassTransformer().to_literal(ctx, python_val, python_type, expected)
+    literal_json = _json_format.MessageToJson(literal.scalar.generic)
+    assert json.loads(literal_json) == python_val
+
+    # Test when python_val is not a dict and not a dataclass
+    python_val = "not a dict or dataclass"
+    with pytest.raises(
+        TypeTransformerFailedError,
+        match="not of type @dataclass, only Dataclasses are supported for user defined datatypes in Flytekit",
+    ):
+        DataclassTransformer().to_literal(ctx, python_val, python_type, expected)
+
+
 @dataclass
 class ArgsAssert(DataClassJSONMixin):
     x: int