Add support FlyteSchema in dataclass

flytekit/core/type_engine.py

@@ -28,7 +28,7 @@
 from flytekit.models import interface as _interface_models
 from flytekit.models import types as _type_models
 from flytekit.models.core import types as _core_types
-from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar
+from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Schema
 from flytekit.models.types import LiteralType, SimpleType
 
 T = typing.TypeVar("T")
@@ -249,10 +249,39 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp
             raise AssertionError(
                 f"Dataclass {python_type} should be decorated with @dataclass_json to be " f"serialized correctly"
             )
+        self._serialize_flyte_type(python_val, python_type)
         return Literal(
             scalar=Scalar(generic=_json_format.Parse(cast(DataClassJsonMixin, python_val).to_json(), _struct.Struct()))
         )
 
+    def _serialize_flyte_type(self, python_val: T, python_type: Type[T]):
+        """
+        If any field inside the dataclass is flyte type, we should use flyte type transformer for that field.
+        """
+        from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer
+
+        for f in dataclasses.fields(python_type):
+            v = python_val.__getattribute__(f.name)
+            if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema):
+                FlyteSchemaTransformer().to_literal(FlyteContext.current_context(), v, f.type, None)
+            elif dataclasses.is_dataclass(f.type):
+                self._serialize_flyte_type(v, f.type)
+
+    def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type["FlyteSchema"]):
+        from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer
+
+        for f in dataclasses.fields(expected_python_type):
+            v = python_val.__getattribute__(f.name)
+            if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema):
+                t = FlyteSchemaTransformer()
+                t.to_python_value(
+                    FlyteContext.current_context(),
+                    Literal(scalar=Scalar(schema=Schema(v.remote_path, t._get_schema_type(f.type)))),
+                    f.type,
+                )
+            elif dataclasses.is_dataclass(f.type):
+                self._deserialize_flyte_type(v, f.type)
+
     def _fix_val_int(self, t: typing.Type, val: typing.Any) -> typing.Any:
         if t == int:
             return int(val)
@@ -295,7 +324,9 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
                 f"Dataclass {expected_python_type} should be decorated with @dataclass_json to be "
                 f"serialized correctly"
             )
+
         dc = cast(DataClassJsonMixin, expected_python_type).from_json(_json_format.MessageToJson(lv.scalar.generic))
+        self._deserialize_flyte_type(dc, expected_python_type)
         return self._fix_dataclass_int(expected_python_type, dc)
 
     def guess_python_type(self, literal_type: LiteralType) -> Type[T]:

flytekit/types/schema/types.py

@@ -4,11 +4,13 @@
 import os
 import typing
 from abc import abstractmethod
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Type
 
 import numpy as _np
+from dataclasses_json import config, dataclass_json
+from marshmallow import fields
 
 from flytekit.core.context_manager import FlyteContext, FlyteContextManager
 from flytekit.core.type_engine import T, TypeEngine, TypeTransformer
@@ -167,7 +169,10 @@ def get_handler(cls, t: Type) -> SchemaHandler:
         return cls._SCHEMA_HANDLERS[t]
 
 
+@dataclass_json
+@dataclass
 class FlyteSchema(object):
+    remote_path: typing.Optional[os.PathLike] = field(default=None, metadata=config(mm_field=fields.String()))
     """
     This is the main schema class that users should use.
     """
@@ -220,7 +225,7 @@ def format(cls) -> SchemaFormat:
     def __init__(
         self,
         local_path: os.PathLike = None,
-        remote_path: str = None,
+        remote_path: os.PathLike = None,
         supported_mode: SchemaOpenMode = SchemaOpenMode.WRITE,
         downloader: typing.Callable[[str, os.PathLike], None] = None,
     ):
@@ -234,10 +239,11 @@ def __init__(
         ):
             raise ValueError("To create a FlyteSchema in write mode, local_path is required")
 
-        if local_path is None:
-            local_path = FlyteContextManager.current_context().file_access.get_random_local_directory()
+        local_path = local_path or FlyteContextManager.current_context().file_access.get_random_local_directory()
         self._local_path = local_path
-        self._remote_path = remote_path
+        # Make this field public, so that the dataclass transformer can set a value for it
+        # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298
+        self.remote_path = remote_path or FlyteContextManager.current_context().file_access.get_random_remote_path()
         self._supported_mode = supported_mode
         # This is a special attribute that indicates if the data was either downloaded or uploaded
         self._downloaded = False
@@ -247,10 +253,6 @@ def __init__(
     def local_path(self) -> os.PathLike:
         return self._local_path
 
-    @property
-    def remote_path(self) -> str:
-        return typing.cast(str, self._remote_path)
-
     @property
     def supported_mode(self) -> SchemaOpenMode:
         return self._supported_mode

tests/flytekit/unit/core/test_type_engine.py

@@ -5,6 +5,7 @@
 from datetime import timedelta
 from enum import Enum
 
+import pandas as pd
 import pytest
 from dataclasses_json import DataClassJsonMixin, dataclass_json
 from flyteidl.core import errors_pb2
@@ -13,6 +14,7 @@
 from marshmallow_enum import LoadDumpOptions
 from marshmallow_jsonschema import JSONSchema
 
+from flytekit import kwtypes
 from flytekit.common.exceptions import user as user_exceptions
 from flytekit.core.context_manager import FlyteContext, FlyteContextManager
 from flytekit.core.type_engine import (
@@ -33,6 +35,7 @@
 from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer
 from flytekit.types.pickle import FlytePickle
 from flytekit.types.pickle.pickle import FlytePickleTransformer
+from flytekit.types.schema import FlyteSchema
 
 
 def test_type_engine():
@@ -657,3 +660,34 @@ def test_dict_to_literal_map_with_wrong_input_type():
     guessed_python_types = {"a": str}
     with pytest.raises(user_exceptions.FlyteTypeException):
         TypeEngine.dict_to_literal_map(ctx, input, guessed_python_types)
+
+
+TestSchema = FlyteSchema[kwtypes(some_str=str)]
+
+
+@dataclass_json
+@dataclass
+class InnerResult:
+    number: int
+    schema: TestSchema
+
+
+@dataclass_json
+@dataclass
+class Result:
+    result: InnerResult
+    schema: TestSchema
+
+
+def test_schema_in_dataclass():
+    schema = TestSchema()
+    df = pd.DataFrame(data={"some_str": ["a", "b", "c"]})
+    schema.open().write(df)
+    o = Result(result=InnerResult(number=1, schema=schema), schema=schema)
+    ctx = FlyteContext.current_context()
+    tf = DataclassTransformer()
+    lt = tf.get_literal_type(Result)
+    lv = tf.to_literal(ctx, o, Result, lt)
+    ot = tf.to_python_value(ctx, lv=lv, expected_python_type=Result)
+
+    assert o == ot

tests/flytekit/unit/core/test_type_hints.py

@@ -1087,6 +1087,35 @@ def wf(x: int) -> Datum:
     assert wf(x=10) == Datum(10, Color.RED)
 
 
+def test_flyte_schema_dataclass():
+    TestSchema = FlyteSchema[kwtypes(some_str=str)]
+
+    @dataclass_json
+    @dataclass
+    class InnerResult:
+        number: int
+        schema: TestSchema
+
+    @dataclass_json
+    @dataclass
+    class Result:
+        result: InnerResult
+        schema: TestSchema
+
+    schema = TestSchema()
+
+    @task
+    def t1(x: int) -> Result:
+
+        return Result(result=InnerResult(number=x, schema=schema), schema=schema)
+
+    @workflow
+    def wf(x: int) -> Result:
+        return t1(x=x)
+
+    assert wf(x=10) == Result(result=InnerResult(number=10, schema=schema), schema=schema)
+
+
 def test_environment():
     @task(environment={"FOO": "foofoo", "BAZ": "baz"})
     def t1(a: int) -> str: