[fix] Validate interface variable names are python-valid (flyteorg#2538)

* Validate interface variable names Signed-off-by: ddl-rliu <[email protected]> * Remove unused import Signed-off-by: ddl-rliu <[email protected]> * Fix lint error Signed-off-by: ddl-rliu <[email protected]> --------- Signed-off-by: ddl-rliu <[email protected]> Signed-off-by: bugra.gedik <[email protected]>
bgedik · Jul 3, 2024 · 8cb7bad · 8cb7bad
1 parent b29b944
commit 8cb7bad
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git 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
@@ -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()