Adding support for ForwardRef

Author: marcromeyn

nemo_run/cli/cli_parser.py

@@ -609,6 +609,7 @@ def __init__(self, strict_mode: bool = True):
             Optional: self.parse_optional,
             Literal: self.parse_literal,
             Path: self.parse_path,
+            'ForwardRef': self.parse_forward_ref,  # Add ForwardRef handling
         }
         self.custom_parsers = {}
         self.strict_mode = strict_mode
@@ -646,6 +647,8 @@ def get_parser(self, annotation: Type) -> Callable[[str, Type], Any]:
             Callable[[str, Type], Any]: The parser function for the given type.
         """
         origin = get_origin(annotation) or annotation
+        if str(origin).startswith('ForwardRef'):
+            return self.parse_forward_ref
         return self.custom_parsers.get(origin) or self.parsers.get(origin) or self.parse_unknown
 
     def parse(self, value: str, annotation: Type) -> Any:
@@ -950,6 +953,23 @@ def parse_path(self, value: str, _: Type) -> Path:
             raise ParseError(value, Path, "Invalid path: contains null character")
         return Path(value.strip("'\" "))
 
+    def parse_forward_ref(self, value: str, annotation) -> Any:
+        """Parse a string value as a ForwardRef type.
+
+        Args:
+            value (str): The string value to parse.
+            annotation: The ForwardRef type annotation.
+
+        Returns:
+            Any: The parsed value.
+
+        Raises:
+            ParseError: If the value cannot be parsed.
+        """
+        # For ForwardRef types, we'll just return the value as is
+        # since the actual type resolution happens later
+        return value
+
     def infer_type(self, value: str) -> Type:
         """Infer the type of a string value.
 

nemo_run/config.py

@@ -97,7 +97,10 @@ def get_type_namespace(typ: Type | Callable) -> str:
     if isinstance(typ, fdl.Buildable):
         typ = typ.__fn_or_cls__
 
-    return f"{module}.{typ.__qualname__}"
+    _name = getattr(typ, "__qualname__", str(typ))
+    if _name.startswith("ForwardRef"):
+        _name = _name.split(".")[-1]
+    return f"{module}.{_name}"
 
 
 def get_underlying_types(type_hint: typing.Any) -> typing.Set[typing.Type]:

test/cli/test_cli_parser.py

@@ -16,7 +16,7 @@
 import sys
 from pathlib import Path
 from test.dummy_factory import DummyModel
-from typing import Any, Dict, List, Literal, Optional, Type, Union
+from typing import Any, Dict, List, Literal, Optional, Type, Union, ForwardRef, TYPE_CHECKING
 
 import pytest
 
@@ -37,6 +37,9 @@
 )
 from nemo_run.config import Config, Partial
 
+if TYPE_CHECKING:
+    from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+
 
 class TestSimpleValueParsing:
     def test_int_parsing(self):
@@ -153,6 +156,23 @@ def func(a: Literal["red", "green", "blue"]):
         assert "Error parsing argument" in str(exc_info.value)
         assert "Expected one of ('red', 'green', 'blue'), got 'yellow'" in str(exc_info.value)
 
+    def test_forward_ref_parsing(self):
+        # Use string annotation to avoid needing the actual class at runtime
+        def func(tokenizer: Optional[ForwardRef("TokenizerSpec")]):
+            pass
+
+        # Test with string value
+        result = parse_cli_args(func, ["tokenizer=tokenizer_spec"])
+        assert result.tokenizer.hidden == 1000
+
+        # Test with None
+        result = parse_cli_args(func, ["tokenizer=None"])
+        assert result.tokenizer is None
+
+        # Test with null (alternative None syntax)
+        result = parse_cli_args(func, ["tokenizer=null"])
+        assert result.tokenizer is None
+
 
 class TestFactoryFunctionParsing:
     def test_simple_factory_function(self):

test/dummy_factory.py

@@ -14,10 +14,13 @@
 # limitations under the License.
 
 from dataclasses import dataclass
-from typing import List
+from typing import List, TYPE_CHECKING, ForwardRef
 
 import nemo_run as run
 
+if TYPE_CHECKING:
+    from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+
 
 @dataclass
 class DummyModel:
@@ -99,6 +102,12 @@ def plugin_list(arg: int = 20) -> List[run.Plugin]:
     ]
 
 
+@run.cli.factory
+@run.autoconvert
+def tokenizer_spec() -> ForwardRef("TokenizerSpec"):
+    return DummyModel(hidden=1000)
+
+
 def dummy_train(dummy_model: DummyModel, dummy_trainer: DummyTrainer): ...