From 10bd84e0bd8b51589f4ba4c67b27d76c4dc8816a Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 26 Oct 2021 00:19:26 +0800 Subject: [PATCH 1/9] schema in dataclass Signed-off-by: Kevin Su --- flytekit/types/schema/types.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 89b6da3a90..2930cb4a31 100644 --- a/flytekit/types/schema/types.py +++ b/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 dataclass_json, config +from marshmallow import fields from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import T, TypeEngine, TypeTransformer @@ -167,7 +169,12 @@ def get_handler(cls, t: Type) -> SchemaHandler: return cls._SCHEMA_HANDLERS[t] +@dataclass_json +@dataclass class FlyteSchema(object): + supported_mode: typing.Optional[str] = field(default=SchemaOpenMode.WRITE, metadata=config(mm_field=fields.String())) + local_path: typing.Optional[str] = field(default=None, metadata=config(mm_field=fields.String())) + remote_path: typing.Optional[str] = field(default=None, metadata=config(mm_field=fields.String())) """ This is the main schema class that users should use. """ @@ -247,14 +254,29 @@ def __init__( def local_path(self) -> os.PathLike: return self._local_path + @local_path.setter + def local_path(self, local_path): + self._local_path = local_path + @property def remote_path(self) -> str: - return typing.cast(str, self._remote_path) + return self._remote_path + + @remote_path.setter + def remote_path(self, remote_path): + self._remote_path = remote_path @property def supported_mode(self) -> SchemaOpenMode: return self._supported_mode + @supported_mode.setter + def supported_mode(self, supported_mode): + self._supported_mode = supported_mode + + def __hash__(self): + return hash(3) + def open( self, dataframe_fmt: type = pandas.DataFrame, override_mode: SchemaOpenMode = None ) -> typing.Union[SchemaReader, SchemaWriter]: From 704eb59099b95ece6727a40ec15856ba5c2a844a Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 26 Oct 2021 02:59:41 +0800 Subject: [PATCH 2/9] Added tests Signed-off-by: Kevin Su --- flytekit/types/schema/types.py | 41 ++++---------------- tests/flytekit/unit/core/test_type_engine.py | 36 +++++++++++++++++ 2 files changed, 43 insertions(+), 34 deletions(-) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 2930cb4a31..15fc209da4 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -243,40 +243,13 @@ def __init__( if local_path is None: local_path = FlyteContextManager.current_context().file_access.get_random_local_directory() - self._local_path = local_path - self._remote_path = remote_path - self._supported_mode = supported_mode + self.local_path = local_path + self.remote_path = 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 self._downloader = downloader - @property - def local_path(self) -> os.PathLike: - return self._local_path - - @local_path.setter - def local_path(self, local_path): - self._local_path = local_path - - @property - def remote_path(self) -> str: - return self._remote_path - - @remote_path.setter - def remote_path(self, remote_path): - self._remote_path = remote_path - - @property - def supported_mode(self) -> SchemaOpenMode: - return self._supported_mode - - @supported_mode.setter - def supported_mode(self, supported_mode): - self._supported_mode = supported_mode - - def __hash__(self): - return hash(3) - def open( self, dataframe_fmt: type = pandas.DataFrame, override_mode: SchemaOpenMode = None ) -> typing.Union[SchemaReader, SchemaWriter]: @@ -291,14 +264,14 @@ def open( So if you have written to a schema and want to re-open it for reading, you can use this mode. A ReadOnly Schema object cannot be opened in write mode. """ - if override_mode and self._supported_mode == SchemaOpenMode.READ and override_mode == SchemaOpenMode.WRITE: + if override_mode and self.supported_mode == SchemaOpenMode.READ and override_mode == SchemaOpenMode.WRITE: raise AssertionError("Readonly schema cannot be opened in write mode!") - mode = override_mode if override_mode else self._supported_mode + mode = override_mode if override_mode else self.supported_mode h = SchemaEngine.get_handler(dataframe_fmt) if not h.handles_remote_io: # The Schema Handler does not manage its own IO, and this it will expect the files are on local file-system - if self._supported_mode == SchemaOpenMode.READ and not self._downloaded: + if self.supported_mode == SchemaOpenMode.READ and not self._downloaded: # Only for readable objects if they are not downloaded already, we should download them # Write objects should already have everything written to self._downloader(self.remote_path, self.local_path) @@ -313,7 +286,7 @@ def open( return h.reader(self.remote_path, self.columns(), self.format()) def as_readonly(self) -> FlyteSchema: - if self._supported_mode == SchemaOpenMode.READ: + if self.supported_mode == SchemaOpenMode.READ: return self s = FlyteSchema.__class_getitem__(self.columns(), self.format())( local_path=self.local_path, diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 3633cb194a..055541ea08 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/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 workflow, task, 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,36 @@ 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 Result: + number: int + schema: TestSchema + + +@task +def t1() -> Result: + schema = TestSchema() + df = pd.DataFrame(data={"some_str": ["a", "b", "c"]}) + schema.open().write(df) + + return Result(number=1, schema=schema) + + +@workflow +def wf() -> Result: + return t1() + + +def test_schema_in_dataclass(): + schema = TestSchema("/tmp") # type: ignore + df = pd.DataFrame(data={"some_str": ["a", "b", "c"]}) + schema.open().write(df) + assert wf().number == 1 + assert "/tmp/flyte" in wf().schema.local_path + assert wf().schema.supported_mode == "w" From b60b3105ad8455742bcb22a38ef6daace05365a6 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 26 Oct 2021 17:06:46 +0800 Subject: [PATCH 3/9] Fixed lint Signed-off-by: Kevin Su --- flytekit/types/schema/types.py | 6 ++++-- tests/flytekit/unit/core/test_type_engine.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 15fc209da4..6c98877419 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -9,7 +9,7 @@ from typing import Type import numpy as _np -from dataclasses_json import dataclass_json, config +from dataclasses_json import config, dataclass_json from marshmallow import fields from flytekit.core.context_manager import FlyteContext, FlyteContextManager @@ -172,7 +172,9 @@ def get_handler(cls, t: Type) -> SchemaHandler: @dataclass_json @dataclass class FlyteSchema(object): - supported_mode: typing.Optional[str] = field(default=SchemaOpenMode.WRITE, metadata=config(mm_field=fields.String())) + supported_mode: typing.Optional[str] = field( + default=SchemaOpenMode.WRITE, metadata=config(mm_field=fields.String()) + ) local_path: typing.Optional[str] = field(default=None, metadata=config(mm_field=fields.String())) remote_path: typing.Optional[str] = field(default=None, metadata=config(mm_field=fields.String())) """ diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 055541ea08..3c237fdc9c 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -14,7 +14,7 @@ from marshmallow_enum import LoadDumpOptions from marshmallow_jsonschema import JSONSchema -from flytekit import workflow, task, kwtypes +from flytekit import kwtypes, task, workflow from flytekit.common.exceptions import user as user_exceptions from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import ( From 4d334759a3e9e8a7604cd6f28e8adfab884f4d66 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 26 Oct 2021 18:00:17 +0800 Subject: [PATCH 4/9] Updated tests Signed-off-by: Kevin Su --- flytekit/types/schema/types.py | 10 ++++------ tests/flytekit/unit/core/test_type_engine.py | 17 +++++++++++++---- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 6c98877419..3cac4f0e2d 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -172,11 +172,9 @@ def get_handler(cls, t: Type) -> SchemaHandler: @dataclass_json @dataclass class FlyteSchema(object): - supported_mode: typing.Optional[str] = field( - default=SchemaOpenMode.WRITE, metadata=config(mm_field=fields.String()) - ) - local_path: typing.Optional[str] = field(default=None, metadata=config(mm_field=fields.String())) - remote_path: typing.Optional[str] = field(default=None, metadata=config(mm_field=fields.String())) + supported_mode: str = field(default=SchemaOpenMode.WRITE, metadata=config(mm_field=fields.String())) + local_path: typing.Optional[os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) + 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. """ @@ -229,7 +227,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, ): diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 3c237fdc9c..a9cdb83862 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -687,9 +687,18 @@ def wf() -> Result: def test_schema_in_dataclass(): - schema = TestSchema("/tmp") # type: ignore + schema = TestSchema() df = pd.DataFrame(data={"some_str": ["a", "b", "c"]}) schema.open().write(df) - assert wf().number == 1 - assert "/tmp/flyte" in wf().schema.local_path - assert wf().schema.supported_mode == "w" + o = Result(number=1, schema=schema) + ctx = FlyteContext.current_context() + tf = DataclassTransformer() + lt = tf.get_literal_type(Result) + gt = tf.guess_python_type(lt) + lv = tf.to_literal(ctx, o, Result, lt) + ot = tf.to_python_value(ctx, lv=lv, expected_python_type=gt) + + assert o.number == ot.number + assert o.schema.local_path == ot.schema.local_path + assert o.schema.remote_path == ot.schema.remote_path + assert o.schema.supported_mode.value == ot.schema.supported_mode From 199f4e58ba1a66a2b27746fa238f1dfa8d58865f Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 27 Oct 2021 01:07:41 +0800 Subject: [PATCH 5/9] Updated tests Signed-off-by: Kevin Su --- tests/flytekit/unit/core/test_type_engine.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index a9cdb83862..bfe5b063ce 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -672,20 +672,6 @@ class Result: schema: TestSchema -@task -def t1() -> Result: - schema = TestSchema() - df = pd.DataFrame(data={"some_str": ["a", "b", "c"]}) - schema.open().write(df) - - return Result(number=1, schema=schema) - - -@workflow -def wf() -> Result: - return t1() - - def test_schema_in_dataclass(): schema = TestSchema() df = pd.DataFrame(data={"some_str": ["a", "b", "c"]}) From 247251fdfe591107a4d82f23533f3e7eab6f1cc2 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 27 Oct 2021 01:16:17 +0800 Subject: [PATCH 6/9] Fixed lint Signed-off-by: Kevin Su --- tests/flytekit/unit/core/test_type_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index bfe5b063ce..6b07f637ef 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -14,7 +14,7 @@ from marshmallow_enum import LoadDumpOptions from marshmallow_jsonschema import JSONSchema -from flytekit import kwtypes, task, workflow +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 ( From 9a1d80cfdc5cd5cc1962117e31d5de8bc0de4139 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 25 Nov 2021 03:51:27 +0800 Subject: [PATCH 7/9] updated Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 31 ++++++++++++++++++++ flytekit/types/schema/types.py | 21 ++++++++----- tests/flytekit/unit/core/test_type_engine.py | 19 +++++++----- tests/flytekit/unit/core/test_type_hints.py | 28 ++++++++++++++++++ 4 files changed, 84 insertions(+), 15 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index c95d6a1576..e970189b91 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -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._flyte_type_to_literal(python_val, python_type) return Literal( scalar=Scalar(generic=_json_format.Parse(cast(DataClassJsonMixin, python_val).to_json(), _struct.Struct())) ) + def _flyte_type_to_literal(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 issubclass(f.type, FlyteSchema): + FlyteSchemaTransformer().to_literal(FlyteContext.current_context(), v, f.type, None) + elif dataclasses.is_dataclass(f.type): + self._flyte_type_to_literal(v, f.type) + + def _flyte_type_to_python_value(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 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._flyte_type_to_python_value(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._flyte_type_to_python_value(dc, expected_python_type) return self._fix_dataclass_int(expected_python_type, dc) def guess_python_type(self, literal_type: LiteralType) -> Type[T]: diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 3cac4f0e2d..c63f1a3a09 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -172,8 +172,6 @@ def get_handler(cls, t: Type) -> SchemaHandler: @dataclass_json @dataclass class FlyteSchema(object): - supported_mode: str = field(default=SchemaOpenMode.WRITE, metadata=config(mm_field=fields.String())) - local_path: typing.Optional[os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) 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. @@ -241,15 +239,24 @@ 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() - self.local_path = local_path - self.remote_path = remote_path - self.supported_mode = supported_mode + local_path = local_path or FlyteContextManager.current_context().file_access.get_random_local_directory() + self._local_path = local_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 self._downloader = downloader + @property + def local_path(self) -> os.PathLike: + return self._local_path + + @property + def supported_mode(self) -> SchemaOpenMode: + return self._supported_mode + def open( self, dataframe_fmt: type = pandas.DataFrame, override_mode: SchemaOpenMode = None ) -> typing.Union[SchemaReader, SchemaWriter]: diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 6b07f637ef..826c8df1fb 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -667,24 +667,27 @@ def test_dict_to_literal_map_with_wrong_input_type(): @dataclass_json @dataclass -class Result: +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(number=1, schema=schema) + o = Result(result=InnerResult(number=1, schema=schema), schema=schema) ctx = FlyteContext.current_context() tf = DataclassTransformer() lt = tf.get_literal_type(Result) - gt = tf.guess_python_type(lt) lv = tf.to_literal(ctx, o, Result, lt) - ot = tf.to_python_value(ctx, lv=lv, expected_python_type=gt) + ot = tf.to_python_value(ctx, lv=lv, expected_python_type=Result) - assert o.number == ot.number - assert o.schema.local_path == ot.schema.local_path - assert o.schema.remote_path == ot.schema.remote_path - assert o.schema.supported_mode.value == ot.schema.supported_mode + assert o == ot diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 02d131ff1c..0ae2566ba7 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1086,6 +1086,34 @@ 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"}) From 304b764de488cfcde2678846fa6934855308053e Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 25 Nov 2021 04:08:45 +0800 Subject: [PATCH 8/9] updated Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 14 +++++++------- flytekit/types/schema/types.py | 8 ++++---- tests/flytekit/unit/core/test_type_hints.py | 1 + 3 files changed, 12 insertions(+), 11 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index e970189b91..45084f46c9 100644 --- a/flytekit/core/type_engine.py +++ b/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,12 +249,12 @@ 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._flyte_type_to_literal(python_val, python_type) + 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 _flyte_type_to_literal(self, python_val: T, python_type: Type[T]): + 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. """ @@ -265,9 +265,9 @@ def _flyte_type_to_literal(self, python_val: T, python_type: Type[T]): if issubclass(f.type, FlyteSchema): FlyteSchemaTransformer().to_literal(FlyteContext.current_context(), v, f.type, None) elif dataclasses.is_dataclass(f.type): - self._flyte_type_to_literal(v, f.type) + self._serialize_flyte_type(v, f.type) - def _flyte_type_to_python_value(self, python_val: T, expected_python_type: Type["FlyteSchema"]): + 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): @@ -280,7 +280,7 @@ def _flyte_type_to_python_value(self, python_val: T, expected_python_type: Type[ f.type, ) elif dataclasses.is_dataclass(f.type): - self._flyte_type_to_python_value(v, 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: @@ -326,7 +326,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: ) dc = cast(DataClassJsonMixin, expected_python_type).from_json(_json_format.MessageToJson(lv.scalar.generic)) - self._flyte_type_to_python_value(dc, expected_python_type) + 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]: diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index c63f1a3a09..cb421f98e3 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -271,14 +271,14 @@ def open( So if you have written to a schema and want to re-open it for reading, you can use this mode. A ReadOnly Schema object cannot be opened in write mode. """ - if override_mode and self.supported_mode == SchemaOpenMode.READ and override_mode == SchemaOpenMode.WRITE: + if override_mode and self._supported_mode == SchemaOpenMode.READ and override_mode == SchemaOpenMode.WRITE: raise AssertionError("Readonly schema cannot be opened in write mode!") - mode = override_mode if override_mode else self.supported_mode + mode = override_mode if override_mode else self._supported_mode h = SchemaEngine.get_handler(dataframe_fmt) if not h.handles_remote_io: # The Schema Handler does not manage its own IO, and this it will expect the files are on local file-system - if self.supported_mode == SchemaOpenMode.READ and not self._downloaded: + if self._supported_mode == SchemaOpenMode.READ and not self._downloaded: # Only for readable objects if they are not downloaded already, we should download them # Write objects should already have everything written to self._downloader(self.remote_path, self.local_path) @@ -293,7 +293,7 @@ def open( return h.reader(self.remote_path, self.columns(), self.format()) def as_readonly(self) -> FlyteSchema: - if self.supported_mode == SchemaOpenMode.READ: + if self._supported_mode == SchemaOpenMode.READ: return self s = FlyteSchema.__class_getitem__(self.columns(), self.format())( local_path=self.local_path, diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 0ae2566ba7..a0559377e7 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1086,6 +1086,7 @@ def wf(x: int) -> Datum: assert wf(x=10) == Datum(10, Color.RED) + def test_flyte_schema_dataclass(): TestSchema = FlyteSchema[kwtypes(some_str=str)] From 33b30a06be51b4d88de896b8ff3b2a8a7dd12d44 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 25 Nov 2021 04:15:33 +0800 Subject: [PATCH 9/9] updated Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 45084f46c9..ff9b257b56 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -262,7 +262,7 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): for f in dataclasses.fields(python_type): v = python_val.__getattribute__(f.name) - if issubclass(f.type, FlyteSchema): + 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) @@ -272,7 +272,7 @@ def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type["Fly for f in dataclasses.fields(expected_python_type): v = python_val.__getattribute__(f.name) - if issubclass(f.type, FlyteSchema): + if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema): t = FlyteSchemaTransformer() t.to_python_value( FlyteContext.current_context(),