From d7cc11ae7ccbcc0dc3fc9395decd7357ecaf2c7e Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 14 Nov 2021 15:08:40 +0800 Subject: [PATCH 1/3] Add support enum in dataclass Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 6 ++++- tests/flytekit/unit/core/test_type_engine.py | 27 ++++++++++++++++++++ tests/flytekit/unit/core/test_type_hints.py | 24 +++++++++++++++++ 3 files changed, 56 insertions(+), 1 deletion(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 925f79c770..0012bf1ce3 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -17,6 +17,7 @@ from google.protobuf.json_format import MessageToDict as _MessageToDict from google.protobuf.json_format import ParseDict as _ParseDict from google.protobuf.struct_pb2 import Struct +from marshmallow_enum import LoadDumpOptions from marshmallow_jsonschema import JSONSchema from flytekit.common.exceptions import user as user_exceptions @@ -226,7 +227,9 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: ) schema = None try: - schema = JSONSchema().dump(cast(DataClassJsonMixin, t).schema()) + s = cast(DataClassJsonMixin, t).schema() + s.fields["y"].load_by = LoadDumpOptions.name + schema = JSONSchema().dump(s) except Exception as e: logger.warn("failed to extract schema for object %s, (will run schemaless) error: %s", str(t), e) @@ -777,6 +780,7 @@ def convert_json_schema_to_python_class(schema: dict, schema_name) -> Type[datac :param schema: dict representing valid JSON schema :param schema_name: dataclass name of return type """ + print(schema) attribute_list = [] for property_key, property_val in schema[schema_name]["properties"].items(): # Handle list diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 6c660b7106..a3166a8250 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -1,6 +1,7 @@ import datetime import os import typing +from ctypes import cast from dataclasses import asdict, dataclass from datetime import timedelta from enum import Enum @@ -10,6 +11,7 @@ from flyteidl.core import errors_pb2 from google.protobuf import json_format as _json_format from google.protobuf import struct_pb2 as _struct +from marshmallow_enum import LoadDumpOptions from marshmallow_jsonschema import JSONSchema from flytekit.common.exceptions import user as user_exceptions @@ -547,6 +549,31 @@ def test_enum_type(): TypeEngine.to_literal_type(UnsupportedEnumValues) +def test_enum_in_dataclass(): + @dataclass_json + @dataclass + class Datum(object): + x: int + y: Color + + def __str__(self): + return self.value + + lt = TypeEngine.to_literal_type(Datum) + schema = Datum.schema() + schema.fields["y"].load_by = LoadDumpOptions.name + assert lt.metadata == JSONSchema().dump(schema) + + transformer = DataclassTransformer() + ctx = FlyteContext.current_context() + datum = Datum(5, Color.RED) + lv = transformer.to_literal(ctx, datum, Datum, lt) + gt = transformer.guess_python_type(lt) + pv = transformer.to_python_value(ctx, lv, expected_python_type=gt) + assert datum.x == pv.x + assert datum.y.value == pv.y + + @pytest.mark.parametrize( "python_value,python_types,expected_literal_map", [ diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index e9f54c829b..02d131ff1c 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -6,6 +6,7 @@ import typing from collections import OrderedDict from dataclasses import dataclass +from enum import Enum import pandas import pytest @@ -1063,6 +1064,29 @@ def wf(x: int, y: int) -> Datum: wf(x=10, y=20) +def test_enum_in_dataclass(): + class Color(Enum): + RED = "red" + GREEN = "green" + BLUE = "blue" + + @dataclass_json + @dataclass + class Datum(object): + x: int + y: Color + + @task + def t1(x: int) -> Datum: + return Datum(x=x, y=Color.RED) + + @workflow + def wf(x: int) -> Datum: + return t1(x=x) + + assert wf(x=10) == Datum(10, Color.RED) + + def test_environment(): @task(environment={"FOO": "foofoo", "BAZ": "baz"}) def t1(a: int) -> str: From 99f6a75b547d38b918ecd8d0a8afeb21298772b5 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 16 Nov 2021 16:21:52 +0800 Subject: [PATCH 2/3] Update test Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 10 +++++++--- tests/flytekit/unit/core/test_type_engine.py | 3 --- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 0012bf1ce3..85fc766d7b 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -17,7 +17,8 @@ from google.protobuf.json_format import MessageToDict as _MessageToDict from google.protobuf.json_format import ParseDict as _ParseDict from google.protobuf.struct_pb2 import Struct -from marshmallow_enum import LoadDumpOptions +from marshmallow import fields +from marshmallow_enum import EnumField, LoadDumpOptions from marshmallow_jsonschema import JSONSchema from flytekit.common.exceptions import user as user_exceptions @@ -228,7 +229,11 @@ def get_literal_type(self, t: Type[T]) -> LiteralType: schema = None try: s = cast(DataClassJsonMixin, t).schema() - s.fields["y"].load_by = LoadDumpOptions.name + for _, v in s.fields.items(): + # marshmallow-jsonschema only supports enums loaded by name. + # https://github.com/fuhrysteve/marshmallow-jsonschema/blob/81eada1a0c42ff67de216923968af0a6b54e5dcb/marshmallow_jsonschema/base.py#L228 + if isinstance(v, EnumField): + v.load_by = LoadDumpOptions.name schema = JSONSchema().dump(s) except Exception as e: logger.warn("failed to extract schema for object %s, (will run schemaless) error: %s", str(t), e) @@ -780,7 +785,6 @@ def convert_json_schema_to_python_class(schema: dict, schema_name) -> Type[datac :param schema: dict representing valid JSON schema :param schema_name: dataclass name of return type """ - print(schema) attribute_list = [] for property_key, property_val in schema[schema_name]["properties"].items(): # Handle list diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index a3166a8250..9b689ce5f1 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -556,9 +556,6 @@ class Datum(object): x: int y: Color - def __str__(self): - return self.value - lt = TypeEngine.to_literal_type(Datum) schema = Datum.schema() schema.fields["y"].load_by = LoadDumpOptions.name From 0cb293aca247ba36129028e143db5ac10ef748a3 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 17 Nov 2021 21:17:28 +0800 Subject: [PATCH 3/3] Fixed lint Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 1 - tests/flytekit/unit/core/test_type_engine.py | 1 - 2 files changed, 2 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 85fc766d7b..c95d6a1576 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -17,7 +17,6 @@ from google.protobuf.json_format import MessageToDict as _MessageToDict from google.protobuf.json_format import ParseDict as _ParseDict from google.protobuf.struct_pb2 import Struct -from marshmallow import fields from marshmallow_enum import EnumField, LoadDumpOptions from marshmallow_jsonschema import JSONSchema diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 9b689ce5f1..3633cb194a 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -1,7 +1,6 @@ import datetime import os import typing -from ctypes import cast from dataclasses import asdict, dataclass from datetime import timedelta from enum import Enum