From 2993693b392a2774e722088091b19a01721b65cb Mon Sep 17 00:00:00 2001 From: ddl-rliu <140021987+ddl-rliu@users.noreply.github.com> Date: Thu, 27 Jun 2024 18:08:14 -0700 Subject: [PATCH] [fix] Validate interface variable names are python-valid (#2538) * Validate interface variable names Signed-off-by: ddl-rliu * Remove unused import Signed-off-by: ddl-rliu * Fix lint error Signed-off-by: ddl-rliu --------- Signed-off-by: ddl-rliu Signed-off-by: Jan Fiedler --- flytekit/core/interface.py | 2 ++ tests/flytekit/unit/core/test_interface.py | 10 ++++++++++ 2 files changed, 12 insertions(+) diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 7f0af5f8e6..7faf541661 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -70,6 +70,8 @@ def __init__( self._inputs: Union[Dict[str, Tuple[Type, Any]], Dict[str, Type]] = {} # type: ignore if inputs: for k, v in inputs.items(): + if not k.isidentifier(): + raise ValueError(f"Input name must be valid Python identifier: {k!r}") if type(v) is tuple and len(cast(Tuple, v)) > 1: self._inputs[k] = v # type: ignore else: diff --git a/tests/flytekit/unit/core/test_interface.py b/tests/flytekit/unit/core/test_interface.py index ecf25fb6e0..e860729a83 100644 --- a/tests/flytekit/unit/core/test_interface.py +++ b/tests/flytekit/unit/core/test_interface.py @@ -320,6 +320,16 @@ def z(a: int, b: str) -> typing.NamedTuple("NT", x_str=str, y_int=int): assert typed_interface.outputs.get("y_int").description == "description for y_int" +def test_init_interface_with_invalid_parameters(): + from flytekit.core.interface import Interface + + with pytest.raises(ValueError, match=r"Input name must be valid Python identifier:"): + _ = Interface({"my.input": int}, {}) + + with pytest.raises(ValueError, match=r"Type names and field names must be valid identifiers:"): + _ = Interface({}, {"my.output": int}) + + def test_parameter_change_to_pickle_type(): ctx = context_manager.FlyteContext.current_context()