Dispatch PEP 604 unions (X | Y) in TypeParser.get_parser (#558)

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

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):