From 631dc40fc4b2d43531303113c0a717c76d2b1a5b Mon Sep 17 00:00:00 2001
From: ddl-rliu <richard.liu@dominodatalab.com>
Date: Thu, 27 Jun 2024 14:47:17 -0700
Subject: [PATCH] Validate interface variable names

Signed-off-by: ddl-rliu <richard.liu@dominodatalab.com>
---
 flytekit/core/interface.py                 |  3 +++
 tests/flytekit/unit/core/test_interface.py | 10 ++++++++++
 2 files changed, 13 insertions(+)

diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py
index 7f0af5f8e6..5ffbbb3ab5 100644
--- a/flytekit/core/interface.py
+++ b/flytekit/core/interface.py
@@ -3,6 +3,7 @@
 import collections
 import copy
 import inspect
+import re
 import typing
 from collections import OrderedDict
 from typing import Any, Dict, Generator, List, Optional, Tuple, Type, TypeVar, Union, cast
@@ -70,6 +71,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 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..33f9e16531 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 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()