From 3ce195a69864d477cad860ff5eabb7f9b0b8a050 Mon Sep 17 00:00:00 2001
From: Srijan Upadhyay <srjnupadhyay@gmail.com>
Date: Wed, 15 Jul 2026 00:53:38 +0530
Subject: [PATCH] Dispatch PEP 604 unions (X | Y) in TypeParser.get_parser

get_origin() returns types.UnionType for PEP 604 unions, not
typing.Union, so annotations like 'str | None' fell through to
'Unsupported type'. Recognize types.UnionType alongside Union at both
dispatch sites; parse_union already handles the members via get_args.
Adds a regression test that fails without this change.

Closes #558

Signed-off-by: Srijan Upadhyay <srjnupadhyay@gmail.com>
---
 nemo_run/cli/cli_parser.py  |  5 +++--
 test/cli/test_cli_parser.py | 23 +++++++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

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