diff --git a/nemo_run/cli/cli_parser.py b/nemo_run/cli/cli_parser.py index 76d8cca1..01e0ac45 100644 --- a/nemo_run/cli/cli_parser.py +++ b/nemo_run/cli/cli_parser.py @@ -21,6 +21,7 @@ import operator import re import sys +import types from enum import Enum from functools import lru_cache from pathlib import Path @@ -685,7 +686,7 @@ def get_parser(self, annotation: Type) -> Callable[[str, Type], Any]: return self.parse_list elif origin in (dict, Dict): return self.parse_dict - elif origin is Union: + elif origin is Union or origin is types.UnionType: return self.parse_union # Add other mappings as needed @@ -704,7 +705,7 @@ def get_parser(self, annotation: Type) -> Callable[[str, Type], Any]: return self.parse_list elif origin is dict or origin is Dict: return self.parse_dict - elif origin is Union: + elif origin is Union or origin is types.UnionType: return self.parse_union # Check for parsers registered for the origin diff --git a/test/cli/test_cli_parser.py b/test/cli/test_cli_parser.py index f8b8afd2..c3fb9a54 100644 --- a/test/cli/test_cli_parser.py +++ b/test/cli/test_cli_parser.py @@ -881,6 +881,29 @@ def func(data: Union[list[str], dict[str, int]]): result = parse_cli_args(func, ["data={'x': 1, 'y': 2}"]) assert result.data == {"x": 1, "y": 2} + def test_modern_pep604_union_type_hints(self): + # PEP 604 unions (X | Y) report get_origin() as types.UnionType, not + # typing.Union, so they previously fell through to "Unsupported type". + # Regression for #558. + if sys.version_info < (3, 10): + pytest.skip("Python 3.10+ required for PEP 604 unions") + + def func(data: list[str] | dict[str, int]): + pass + + result = parse_cli_args(func, ["data=['a', 'b', 'c']"]) + assert result.data == ["a", "b", "c"] + result = parse_cli_args(func, ["data={'x': 1, 'y': 2}"]) + assert result.data == {"x": 1, "y": 2} + + def func_optional(name: str | None): + pass + + result = parse_cli_args(func_optional, ["name=hello"]) + assert result.name == "hello" + result = parse_cli_args(func_optional, ["name=None"]) + assert result.name is None + def test_modern_type_parsing_errors(self): # Skip test if running on Python < 3.9 if sys.version_info < (3, 9):