From dfba3d822ccd409e300e9c98aaee18a5bd689d0b Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 30 Aug 2022 02:50:29 +0800 Subject: [PATCH 1/4] Add literal type to union literal Signed-off-by: Kevin Su --- flytekit/clis/sdk_in_container/run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index 95533fb4d5..e54d5930fc 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -27,7 +27,7 @@ from flytekit.core.workflow import PythonFunctionWorkflow, WorkflowBase from flytekit.models import literals from flytekit.models.interface import Variable -from flytekit.models.literals import Blob, BlobMetadata, Primitive +from flytekit.models.literals import Blob, BlobMetadata, Primitive, Union from flytekit.models.types import LiteralType, SimpleType from flytekit.remote.executions import FlyteWorkflowExecution from flytekit.tools import module_loader, script_mode @@ -265,7 +265,7 @@ def convert_to_union( python_val = converter._click_type.convert(value, param, ctx) literal = converter.convert_to_literal(ctx, param, python_val) self._python_type = python_type - return literal + return Literal(scalar=Scalar(union=Union(literal, variant))) except (Exception or AttributeError) as e: logging.debug(f"Failed to convert python type {python_type} to literal type {variant}", e) raise ValueError(f"Failed to convert python type {self._python_type} to literal type {lt}") From 602780ea501449ca6e2cb7750f5841b6a4b4b641 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 30 Aug 2022 13:43:23 +0800 Subject: [PATCH 2/4] fix test Signed-off-by: Kevin Su --- flytekit/clis/sdk_in_container/run.py | 1 - 1 file changed, 1 deletion(-) diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index e54d5930fc..6365606411 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -264,7 +264,6 @@ def convert_to_union( # and then use flyte converter to convert it to literal. python_val = converter._click_type.convert(value, param, ctx) literal = converter.convert_to_literal(ctx, param, python_val) - self._python_type = python_type return Literal(scalar=Scalar(union=Union(literal, variant))) except (Exception or AttributeError) as e: logging.debug(f"Failed to convert python type {python_type} to literal type {variant}", e) From 1187ddb8af37d48b181c527ab5ba8828e7e135a7 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 31 Aug 2022 01:55:20 +0800 Subject: [PATCH 3/4] Add tests Signed-off-by: Kevin Su --- tests/flytekit/unit/cli/pyflyte/test_run.py | 41 ++++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/tests/flytekit/unit/cli/pyflyte/test_run.py b/tests/flytekit/unit/cli/pyflyte/test_run.py index 34f28f00b6..2e27d18533 100644 --- a/tests/flytekit/unit/cli/pyflyte/test_run.py +++ b/tests/flytekit/unit/cli/pyflyte/test_run.py @@ -1,10 +1,15 @@ +import functools import os import pathlib +import typing +from enum import Enum +import click import mock import pytest from click.testing import CliRunner +from flytekit import FlyteContextManager from flytekit.clis.sdk_in_container import pyflyte from flytekit.clis.sdk_in_container.constants import CTX_CONFIG_FILE from flytekit.clis.sdk_in_container.helpers import FLYTE_REMOTE_INSTANCE_KEY @@ -12,11 +17,14 @@ REMOTE_FLAG_KEY, RUN_LEVEL_PARAMS_KEY, FileParamType, + FlyteLiteralConverter, get_entities_in_file, run_command, ) -from flytekit.configuration import Image, ImageConfig +from flytekit.configuration import Config, Image, ImageConfig from flytekit.core.task import task +from flytekit.core.type_engine import TypeEngine +from flytekit.remote import FlyteRemote WORKFLOW_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "workflow.py") DIR_NAME = os.path.dirname(os.path.realpath(__file__)) @@ -255,3 +263,34 @@ def test_file_param(): assert l.local r = FileParamType().convert("https://tmp/file", m, m) assert r.local is False + + +class Color(Enum): + RED = "red" + GREEN = "green" + BLUE = "blue" + + +@pytest.mark.parametrize( + "python_type, python_value", + [ + (typing.Union[typing.List[int], str, Color], "flyte"), + (typing.Union[typing.List[int], str, Color], "red"), + (typing.Union[typing.List[int], str, Color], [1, 2, 3]), + (typing.List[int], [1, 2, 3]), + (typing.Dict[str, int], {"flyte": 2}), + ], +) +def test_literal_converter(python_type, python_value): + get_upload_url_fn = functools.partial( + FlyteRemote(Config.auto()).client.get_upload_signed_url, project="p", domain="d" + ) + click_ctx = click.Context(click.Command("test_command"), obj={"remote": True}) + ctx = FlyteContextManager.current_context() + lt = TypeEngine.to_literal_type(python_type) + + lc = FlyteLiteralConverter( + click_ctx, ctx, literal_type=lt, python_type=python_type, get_upload_url_fn=get_upload_url_fn + ) + + assert lc.convert(click_ctx, ctx, python_value) == TypeEngine.to_literal(ctx, python_value, python_type, lt) From 304dd94deb035767dc3ca8723626efa465746f98 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 1 Sep 2022 00:29:47 +0800 Subject: [PATCH 4/4] more tests Signed-off-by: Kevin Su --- tests/flytekit/unit/cli/pyflyte/test_run.py | 23 +++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/flytekit/unit/cli/pyflyte/test_run.py b/tests/flytekit/unit/cli/pyflyte/test_run.py index 2e27d18533..fd5f5acc14 100644 --- a/tests/flytekit/unit/cli/pyflyte/test_run.py +++ b/tests/flytekit/unit/cli/pyflyte/test_run.py @@ -24,6 +24,7 @@ from flytekit.configuration import Config, Image, ImageConfig from flytekit.core.task import task from flytekit.core.type_engine import TypeEngine +from flytekit.models.types import SimpleType from flytekit.remote import FlyteRemote WORKFLOW_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "workflow.py") @@ -294,3 +295,25 @@ def test_literal_converter(python_type, python_value): ) assert lc.convert(click_ctx, ctx, python_value) == TypeEngine.to_literal(ctx, python_value, python_type, lt) + + +def test_enum_converter(): + pt = typing.Union[str, Color] + + get_upload_url_fn = functools.partial(FlyteRemote(Config.auto()).client.get_upload_signed_url) + click_ctx = click.Context(click.Command("test_command"), obj={"remote": True}) + ctx = FlyteContextManager.current_context() + lt = TypeEngine.to_literal_type(pt) + lc = FlyteLiteralConverter(click_ctx, ctx, literal_type=lt, python_type=pt, get_upload_url_fn=get_upload_url_fn) + union_lt = lc.convert(click_ctx, ctx, "red").scalar.union + + assert union_lt.stored_type.simple == SimpleType.STRING + assert union_lt.stored_type.enum_type is None + + pt = typing.Union[Color, str] + lt = TypeEngine.to_literal_type(typing.Union[Color, str]) + lc = FlyteLiteralConverter(click_ctx, ctx, literal_type=lt, python_type=pt, get_upload_url_fn=get_upload_url_fn) + union_lt = lc.convert(click_ctx, ctx, "red").scalar.union + + assert union_lt.stored_type.simple is None + assert union_lt.stored_type.enum_type.values == ["red", "green", "blue"]