Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 21 additions & 7 deletions nemo_run/cli/cli_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -869,8 +869,13 @@ def parse_list(self, value: str, annotation: Type[List]) -> List:
parsed = ast.literal_eval(value)
if not isinstance(parsed, list):
raise ValueError("Not a list")
elem_type = get_args(annotation)[0]
return [self.parse(str(item), elem_type) for item in parsed]
type_args = get_args(annotation)
if type_args:
elem_type = type_args[0]
return [self.parse(str(item), elem_type) for item in parsed]
# Unparameterized annotations (e.g. `list`, `List`, `Optional[list]`
# resolved to a bare list) have no element type to coerce to.
return parsed
except Exception as e:
raise ListParseError(value, List, f"Invalid list: {str(e)}")

Expand All @@ -891,11 +896,16 @@ def parse_dict(self, value: str, annotation: Type[Dict]) -> Dict:
parsed = ast.literal_eval(value)
if not isinstance(parsed, dict):
raise ValueError("Not a dict")
key_type, val_type = get_args(annotation)
return {
self.parse(str(k), key_type): self.parse(str(v), val_type)
for k, v in parsed.items()
}
type_args = get_args(annotation)
if type_args:
key_type, val_type = type_args
return {
self.parse(str(k), key_type): self.parse(str(v), val_type)
for k, v in parsed.items()
}
# Unparameterized annotations (e.g. `dict`, `Dict`, `Optional[dict]`
# resolved to a bare dict) have no key/value types to coerce to.
return parsed
except Exception as e:
raise DictParseError(value, Dict, f"Invalid dict: {str(e)}")

Expand Down Expand Up @@ -1469,6 +1479,10 @@ def _maybe_resolve_annotation(fn: Callable, arg_name: str, annotation: Any) -> A
# Case 3: Annotation is a generic type (e.g., Optional, List, Union)
elif (origin := get_origin(annotation)) is not None:
args = get_args(annotation)
if not args:
# Unparameterized generics (e.g. bare `List`, `Dict`) have no
# arguments to resolve and are handled downstream as-is.
return annotation
resolved_args = tuple(_maybe_resolve_annotation(fn, arg_name, arg) for arg in args)
if origin is list:
return List[resolved_args[0]]
Expand Down
39 changes: 39 additions & 0 deletions test/cli/test_cli_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,33 @@ def func(a: List[List[int]]):

assert parse_cli_args(func, ["a=[[1, 2], [3, 4]]"]).a == [[1, 2], [3, 4]]

def test_unparameterized_list_parsing(self):
def func(a: list, b: List = None, c: Optional[list] = None):
pass

assert parse_cli_args(func, ["a=[1, 2, 3]"]).a == [1, 2, 3]
assert parse_cli_args(func, ["a=[]"]).a == []
assert parse_cli_args(func, ["b=[1, 2]"]).b == [1, 2]
assert parse_cli_args(func, ["c=[1, 2]"]).c == [1, 2]

def test_unparameterized_dict_parsing(self):
def func(a: dict, b: Dict = None, c: Optional[dict] = None):
pass

assert parse_cli_args(func, ["a={'x': 1}"]).a == {"x": 1}
assert parse_cli_args(func, ["a={}"]).a == {}
assert parse_cli_args(func, ["b={'x': 1}"]).b == {"x": 1}
assert parse_cli_args(func, ["c={'x': 1}"]).c == {"x": 1}

def test_union_with_list_not_misparsed_as_string(self):
def func(a: Union[list, str] = None, b: Union[dict, str] = None):
pass

assert parse_cli_args(func, ["a=[1, 2]"]).a == [1, 2]
assert parse_cli_args(func, ["a=hello"]).a == "hello"
assert parse_cli_args(func, ["b={'x': 1}"]).b == {"x": 1}
assert parse_cli_args(func, ["b=hello"]).b == "hello"

def test_dict_parsing(self):
def func(a: Dict[str, int]):
pass
Expand Down Expand Up @@ -506,6 +533,18 @@ def test_parse_dict(self):
with pytest.raises(ParseError, match="Failed to parse"):
parse_value('{"a": 1, "b": "two"}', Dict[str, int])

def test_parse_unparameterized_list(self):
assert parse_value("[1, 2, 3]", list) == [1, 2, 3]
assert parse_value("[1, 2, 3]", List) == [1, 2, 3]
assert parse_value("[1, 2, 3]", Optional[list]) == [1, 2, 3]
assert parse_value("None", Optional[list]) is None

def test_parse_unparameterized_dict(self):
assert parse_value('{"a": 1}', dict) == {"a": 1}
assert parse_value('{"a": 1}', Dict) == {"a": 1}
assert parse_value('{"a": 1}', Optional[dict]) == {"a": 1}
assert parse_value("None", Optional[dict]) is None

def test_parse_union(self):
assert parse_value("123", Union[int, str]) == 123
assert parse_value("hello", Union[int, str]) == "hello"
Expand Down