Skip to content

Commit

Permalink
Validate interface variable names
Browse files Browse the repository at this point in the history
Signed-off-by: ddl-rliu <[email protected]>
  • Loading branch information
ddl-rliu committed Jun 27, 2024
1 parent bd4c708 commit bb0f43d
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 0 deletions.
3 changes: 3 additions & 0 deletions flytekit/core/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 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:
Expand Down
10 changes: 10 additions & 0 deletions tests/flytekit/unit/core/test_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down

0 comments on commit bb0f43d

Please sign in to comment.