Support mashumaro DataClassORJSONMixin (#2080)

* support DataClassJSONMixin

Signed-off-by: Quinten Roets <quinten.roets@flawlessai.com>

* make union syntax compatible with python3.8

Signed-off-by: Quinten Roets <quinten.roets@flawlessai.com>

* add datetime attribute

Signed-off-by: Quinten Roets <quinten.roets@flawlessai.com>

* centralize serializable checking

Signed-off-by: Quinten Roets <quinten.roets@flawlessai.com>

* Incorporate feedback

Signed-off-by: Eduardo Apolinario <eapolinario@users.noreply.github.com>

---------

Signed-off-by: Quinten Roets <quinten.roets@flawlessai.com>
Signed-off-by: Eduardo Apolinario <eapolinario@users.noreply.github.com>
Co-authored-by: Eduardo Apolinario <eapolinario@users.noreply.github.com>

Author: quinten-flwls

dev-requirements.in

@@ -45,3 +45,5 @@ pandas
 scikit-learn
 types-requests
 prometheus-client
+
+orjson

flytekit/core/type_engine.py

@@ -325,6 +325,13 @@ class Test(DataClassJsonMixin):
 
     def __init__(self):
         super().__init__("Object-Dataclass-Transformer", object)
+        self._serializable_classes = [DataClassJSONMixin, DataClassJsonMixin]
+        try:
+            from mashumaro.mixins.orjson import DataClassORJSONMixin
+
+            self._serializable_classes.append(DataClassORJSONMixin)
+        except ModuleNotFoundError:
+            pass
 
     def assert_type(self, expected_type: Type[DataClassJsonMixin], v: T):
         # Skip iterating all attributes in the dataclass if the type of v already matches the expected_type
@@ -417,7 +424,7 @@ def get_literal_type(self, t: Type[T]) -> LiteralType:
                 f"Type {t} cannot be parsed."
             )
 
-        if not issubclass(t, DataClassJsonMixin) and not issubclass(t, DataClassJSONMixin):
+        if not self.is_serializable_class(t):
             raise AssertionError(
                 f"Dataclass {t} should be decorated with @dataclass_json or mixin with DataClassJSONMixin to be "
                 f"serialized correctly"
@@ -465,6 +472,9 @@ def get_literal_type(self, t: Type[T]) -> LiteralType:
 
         return _type_models.LiteralType(simple=_type_models.SimpleType.STRUCT, metadata=schema, structure=ts)
 
+    def is_serializable_class(self, class_: Type[T]) -> bool:
+        return any(issubclass(class_, serializable_class) for serializable_class in self._serializable_classes)
+
     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)
@@ -475,9 +485,7 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp
                 f"{type(python_val)} is not of type @dataclass, only Dataclasses are supported for "
                 f"user defined datatypes in Flytekit"
             )
-        if not issubclass(type(python_val), DataClassJsonMixin) and not issubclass(
-            type(python_val), DataClassJSONMixin
-        ):
+        if not self.is_serializable_class(type(python_val)):
             raise TypeTransformerFailedError(
                 f"Dataclass {python_type} should be decorated with @dataclass_json or inherit DataClassJSONMixin to be "
                 f"serialized correctly"
@@ -730,9 +738,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
                 f"{expected_python_type} is not of type @dataclass, only Dataclasses are supported for "
                 "user defined datatypes in Flytekit"
             )
-        if not issubclass(expected_python_type, DataClassJsonMixin) and not issubclass(
-            expected_python_type, DataClassJSONMixin
-        ):
+        if not self.is_serializable_class(expected_python_type):
             raise TypeTransformerFailedError(
                 f"Dataclass {expected_python_type} should be decorated with @dataclass_json or mixin with DataClassJSONMixin to be "
                 f"serialized correctly"

tests/flytekit/unit/core/test_type_engine.py

@@ -22,6 +22,7 @@
 from marshmallow_enum import LoadDumpOptions
 from marshmallow_jsonschema import JSONSchema
 from mashumaro.mixins.json import DataClassJSONMixin
+from mashumaro.mixins.orjson import DataClassORJSONMixin
 from typing_extensions import Annotated, get_args, get_origin
 
 from flytekit import kwtypes
@@ -2366,6 +2367,10 @@ def test_DataclassTransformer_get_literal_type():
     class MyDataClassMashumaro(DataClassJsonMixin):
         x: int
 
+    @dataclass
+    class MyDataClassMashumaroORJSON(DataClassJsonMixin):
+        x: int
+
     @dataclass_json
     @dataclass
     class MyDataClass:
@@ -2379,6 +2384,9 @@ class MyDataClass:
     literal_type = de.get_literal_type(MyDataClassMashumaro)
     assert literal_type is not None
 
+    literal_type = de.get_literal_type(MyDataClassMashumaroORJSON)
+    assert literal_type is not None
+
     invalid_json_str = "{ unbalanced_braces"
     with pytest.raises(Exception):
         Literal(scalar=Scalar(generic=_json_format.Parse(invalid_json_str, _struct.Struct())))
@@ -2389,6 +2397,10 @@ def test_DataclassTransformer_to_literal():
     class MyDataClassMashumaro(DataClassJsonMixin):
         x: int
 
+    @dataclass
+    class MyDataClassMashumaroORJSON(DataClassORJSONMixin):
+        x: int
+
     @dataclass_json
     @dataclass
     class MyDataClass:
@@ -2398,12 +2410,19 @@ class MyDataClass:
     ctx = FlyteContext.current_context()
 
     my_dat_class_mashumaro = MyDataClassMashumaro(5)
+    my_dat_class_mashumaro_orjson = MyDataClassMashumaroORJSON(5)
     my_data_class = MyDataClass(5)
 
     lv_mashumaro = transformer.to_literal(ctx, my_dat_class_mashumaro, MyDataClassMashumaro, MyDataClassMashumaro)
     assert lv_mashumaro is not None
     assert lv_mashumaro.scalar.generic["x"] == 5
 
+    lv_mashumaro_orjson = transformer.to_literal(
+        ctx, my_dat_class_mashumaro_orjson, MyDataClassMashumaroORJSON, MyDataClassMashumaroORJSON
+    )
+    assert lv_mashumaro_orjson is not None
+    assert lv_mashumaro_orjson.scalar.generic["x"] == 5
+
     lv = transformer.to_literal(ctx, my_data_class, MyDataClass, MyDataClass)
     assert lv is not None
     assert lv.scalar.generic["x"] == 5
@@ -2414,6 +2433,10 @@ def test_DataclassTransformer_to_python_value():
     class MyDataClassMashumaro(DataClassJsonMixin):
         x: int
 
+    @dataclass
+    class MyDataClassMashumaroORJSON(DataClassORJSONMixin):
+        x: int
+
     @dataclass_json
     @dataclass
     class MyDataClass:
@@ -2432,8 +2455,18 @@ class MyDataClass:
     assert isinstance(result, MyDataClassMashumaro)
     assert result.x == 5
 
+    result = de.to_python_value(FlyteContext.current_context(), mock_literal, MyDataClassMashumaroORJSON)
+    assert isinstance(result, MyDataClassMashumaroORJSON)
+    assert result.x == 5
+
 
 def test_DataclassTransformer_guess_python_type():
+    @dataclass
+    class DatumMashumaroORJSON(DataClassORJSONMixin):
+        x: int
+        y: Color
+        z: datetime.datetime
+
     @dataclass
     class DatumMashumaro(DataClassJSONMixin):
         x: int
@@ -2464,6 +2497,16 @@ class Datum(DataClassJSONMixin):
     assert datum_mashumaro.x == pv.x
     assert datum_mashumaro.y.value == pv.y
 
+    lt = TypeEngine.to_literal_type(DatumMashumaroORJSON)
+    now = datetime.datetime.now()
+    datum_mashumaro_orjson = DatumMashumaroORJSON(5, Color.RED, now)
+    lv = transformer.to_literal(ctx, datum_mashumaro_orjson, DatumMashumaroORJSON, lt)
+    gt = transformer.guess_python_type(lt)
+    pv = transformer.to_python_value(ctx, lv, expected_python_type=gt)
+    assert datum_mashumaro_orjson.x == pv.x
+    assert datum_mashumaro_orjson.y.value == pv.y
+    assert datum_mashumaro_orjson.z.isoformat() == pv.z
+
 
 def test_ListTransformer_get_sub_type():
     assert ListTransformer.get_sub_type_or_none(typing.List[str]) is str