From 1d2e9b9a0964e9a144e4bfa389eddc27346d1867 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Thu, 12 May 2022 15:41:56 -0700 Subject: [PATCH] allow users to just have one --pkgs switch, but have commas in there indicating separate packages Signed-off-by: Yee Hing Tong --- flytekit/clis/sdk_in_container/pyflyte.py | 12 ++++++++++-- tests/flytekit/unit/cli/pyflyte/test_package.py | 5 +++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/flytekit/clis/sdk_in_container/pyflyte.py b/flytekit/clis/sdk_in_container/pyflyte.py index f27128d108..edbd5da79d 100644 --- a/flytekit/clis/sdk_in_container/pyflyte.py +++ b/flytekit/clis/sdk_in_container/pyflyte.py @@ -8,15 +8,22 @@ from flytekit.clis.sdk_in_container.run import run from flytekit.clis.sdk_in_container.serialize import serialize from flytekit.configuration.internal import LocalSDK +from flytekit.loggers import cli_logger def validate_package(ctx, param, values): + pkgs = [] for val in values: if "/" in val or "-" in val or "\\" in val: raise click.BadParameter( f"Illegal package value {val} for parameter: {param}. Expected for the form [a.b.c]" ) - return values + elif "," in val: + pkgs.extend(val.split(",")) + else: + pkgs.append(val) + cli_logger.debug(f"Using packages: {pkgs}") + return pkgs @click.group("pyflyte", invoke_without_command=True) @@ -26,7 +33,8 @@ def validate_package(ctx, param, values): required=False, multiple=True, callback=validate_package, - help="Dot separated python packages to operate on. Multiple may be specified Please note that this " + help="Dot-delineated python packages to operate on. Multiple may be specified (can use commas, or specify the " + "switch multiple times. Please note that this " "option will override the option specified in the configuration file, or environment variable", ) @click.option( diff --git a/tests/flytekit/unit/cli/pyflyte/test_package.py b/tests/flytekit/unit/cli/pyflyte/test_package.py index 74b3f7ff6f..d31b760608 100644 --- a/tests/flytekit/unit/cli/pyflyte/test_package.py +++ b/tests/flytekit/unit/cli/pyflyte/test_package.py @@ -109,3 +109,8 @@ def test_package(): ) assert result.exit_code == 1 assert result.output is not None + + +def test_pkgs(): + pp = pyflyte.validate_package(None, None, ["a.b", "a.c,b.a", "cc.a"]) + assert pp == ["a.b", "a.c", "b.a", "cc.a"]