Adding support for ForwardRef in CLI (#176)

* Adding support for ForwardRef

Signed-off-by: Marc Romeyn <[email protected]>

* Ruff formatting

Signed-off-by: Marc Romeyn <[email protected]>

* Adding support for ForwardRef

Signed-off-by: Marc Romeyn <[email protected]>

* Fixing ruff issues

Signed-off-by: Marc Romeyn <[email protected]>

* Add extra check for ForwardRef

Signed-off-by: Marc Romeyn <[email protected]>

* Trying to fix failing test

Signed-off-by: Marc Romeyn <[email protected]>

* Trying to fix failing test

Signed-off-by: Marc Romeyn <[email protected]>

* Fix failing test

Signed-off-by: Marc Romeyn <[email protected]>

* Fix linting issues

Signed-off-by: Marc Romeyn <[email protected]>

* Adding support for Optional[ForwardRef(..)]

Signed-off-by: Marc Romeyn <[email protected]>

* Put back normal pyproject.toml

Signed-off-by: Marc Romeyn <[email protected]>

* Fix bug

Signed-off-by: Marc Romeijn <[email protected]>

* Trying to fix failing test

Signed-off-by: Marc Romeyn <[email protected]>

---------

Signed-off-by: Marc Romeyn <[email protected]>
Signed-off-by: Marc Romeyn <[email protected]>
Signed-off-by: Marc Romeijn <[email protected]>

Author: marcromeyn

nemo_run/cli/cli_parser.py

@@ -36,6 +36,10 @@
     Union,
     get_args,
     get_origin,
+    ForwardRef,
+    Set,
+    Tuple,
+    FrozenSet,
 )
 
 import fiddle as fdl
@@ -609,6 +613,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,
         }
         self.custom_parsers = {}
         self.strict_mode = strict_mode
@@ -646,6 +651,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 +957,9 @@ 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:
+        return value
+
     def infer_type(self, value: str) -> Type:
         """Infer the type of a string value.
 
@@ -980,7 +990,9 @@ def parse_value(value: str, annotation: Type = None) -> Any:
 
 @cli_exception_handler
 def parse_cli_args(
-    fn: Callable, args: List[str], output_type: Type[TypeVar("OutputT", Partial, Config)] = Partial
+    fn: Callable,
+    args: List[str],
+    output_type: Type[TypeVar("OutputT", Partial, Config)] = Partial,
 ) -> TypeVar("OutputT", Partial, Config):
     """Parse command-line arguments and apply them to a function or class.
 
@@ -1109,6 +1121,10 @@ def dummy_model_config():
             annotation = param.annotation
         logger.debug(f"Parsing value {value} as {annotation}")
 
+        annotation = _maybe_resolve_annotation(
+            getattr(nested, "__fn_or_cls__", nested), arg_name, annotation
+        )
+
         if annotation:
             try:
                 parsed_value = parse_factory(fn, arg_name, annotation, value)
@@ -1132,7 +1148,9 @@ def dummy_model_config():
             else:
                 if not hasattr(nested, arg_name):
                     raise UndefinedVariableError(
-                        f"Cannot use '{op.value}' on undefined variable", arg, {"key": key}
+                        f"Cannot use '{op.value}' on undefined variable",
+                        arg,
+                        {"key": key},
                     )
                 setattr(
                     nested,
@@ -1277,7 +1295,9 @@ def _args_to_kwargs(fn: Callable, args: List[str]) -> List[str]:
         for arg in args:
             if "=" not in arg:
                 raise ArgumentParsingError(
-                    "Positional argument found after keyword argument", arg, {"position": len(args)}
+                    "Positional argument found after keyword argument",
+                    arg,
+                    {"position": len(args)},
                 )
 
         return args
@@ -1304,7 +1324,9 @@ def _args_to_kwargs(fn: Callable, args: List[str]) -> List[str]:
                 positional_count += 1
             else:
                 raise ArgumentParsingError(
-                    "Too many positional arguments", arg, {"max_positional": len(params)}
+                    "Too many positional arguments",
+                    arg,
+                    {"max_positional": len(params)},
                 )
 
     return updated_args
@@ -1334,3 +1356,92 @@ def parse_attribute(attr, nested):
                 ) from e
 
     return result
+
+
+def _maybe_resolve_annotation(fn: Callable, arg_name: str, annotation: Any) -> Any:
+    """Internal function to resolve an annotation to its actual type.
+
+    This function handles string annotations, ForwardRef, and generic types (e.g., Optional, List)
+    by resolving string annotations within them, using TYPE_CHECKING blocks and their imports.
+
+    Args:
+        fn (Callable): The function containing the annotation
+        arg_name (str): The name of the parameter with the annotation
+        annotation (Any): The annotation to resolve (string, ForwardRef, or type)
+
+    Returns:
+        Any: The resolved type, or the original annotation if resolution fails
+    """
+    # Case 1: Annotation is a string
+    if isinstance(annotation, str):
+        resolved = _resolve_type_checking_annotation(fn, annotation)
+        return resolved if resolved != annotation else annotation
+
+    # Case 2: Annotation is a ForwardRef
+    elif isinstance(annotation, ForwardRef):
+        return _resolve_type_checking_annotation(fn, annotation.__forward_arg__)
+
+    # Case 3: Annotation is a generic type (e.g., Optional, List, Union)
+    elif (origin := get_origin(annotation)) is not None:
+        args = get_args(annotation)
+        resolved_args = tuple(_maybe_resolve_annotation(fn, arg_name, arg) for arg in args)
+        if origin is list:
+            return List[resolved_args[0]]
+        elif origin is dict:
+            return Dict[resolved_args[0], resolved_args[1]]
+        elif origin is tuple:
+            return Tuple[resolved_args]
+        elif origin is set:
+            return Set[resolved_args[0]]
+        elif origin is frozenset:
+            return FrozenSet[resolved_args[0]]
+        elif origin is Union:
+            return Union[resolved_args]
+        else:
+            return annotation  # Unhandled generic types return as-is
+
+    # Case 4: Annotation is a non-generic type (e.g., int, str)
+    else:
+        return annotation
+
+
+def _resolve_type_checking_annotation(fn: Callable, annotation: str) -> Any:
+    """Helper function to resolve a string annotation to its actual type using TYPE_CHECKING imports."""
+    if hasattr(fn, "__fn_or_cls__"):
+        fn = fn.__fn_or_cls__
+
+    try:
+        source_file = inspect.getsourcefile(fn)
+        if not source_file:
+            return annotation
+        with open(source_file, "r") as f:
+            source = f.read()
+        tree = ast.parse(source)
+        type_checking_imports = {}
+        for node in ast.walk(tree):
+            if (
+                isinstance(node, ast.If)
+                and isinstance(node.test, ast.Name)
+                and node.test.id == "TYPE_CHECKING"
+            ):
+                for stmt in node.body:
+                    if isinstance(stmt, (ast.Import, ast.ImportFrom)):
+                        if isinstance(stmt, ast.Import):
+                            for name in stmt.names:
+                                type_checking_imports[name.asname or name.name] = name.name
+                        else:  # ImportFrom
+                            module = stmt.module or ""
+                            for name in stmt.names:
+                                full_name = f"{module}.{name.name}" if module else name.name
+                                type_checking_imports[name.asname or name.name] = full_name
+        if annotation in type_checking_imports:
+            try:
+                full_path = type_checking_imports[annotation]
+                module_name, type_name = full_path.rsplit(".", 1)
+                module = importlib.import_module(module_name)
+                return getattr(module, type_name)
+            except (ImportError, AttributeError):
+                pass
+    except Exception:
+        pass
+    return annotation

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_api.py

@@ -17,7 +17,7 @@
 import sys
 from configparser import ConfigParser
 from dataclasses import dataclass, field
-from typing import Annotated, List, Optional, Union
+from typing import Annotated, List, Optional, Union, TYPE_CHECKING
 from unittest.mock import Mock, patch
 
 import fiddle as fdl
@@ -46,6 +46,10 @@
 from test.dummy_factory import DummyModel, dummy_entrypoint
 import nemo_run.cli.cli_parser  # Import the module to mock its function
 
+if TYPE_CHECKING:
+    from test.dummy_type import RealType
+
+
 _RUN_FACTORIES_ENTRYPOINT: str = """
 [nemo_run.cli]
 dummy = test.dummy_factory
@@ -519,6 +523,42 @@ def test_factory_for_entrypoint(self):
         cfg = run.cli.resolve_factory(dummy_entrypoint, "dummy_recipe")()
         assert cfg.dummy.hidden == 2000
 
+    def test_forward_ref_with_real_type_factory(self):
+        """Test that ForwardRef works when factory is registered for the actual type."""
+
+        # Function that uses ForwardRef to the module-level RealType class
+        def func(param: Optional["RealType"] = None):
+            pass
+
+        from test.dummy_type import RealType as _RealType
+
+        # Register the factory in the module's global namespace
+        # The factory returns a RealType instance with a specific value
+        @run.cli.factory
+        @run.autoconvert
+        def real_type_factory() -> _RealType:
+            return _RealType(value=100)
+
+        @run.cli.factory(target=func, target_arg="param")
+        @run.autoconvert
+        def other_factory() -> _RealType:
+            return _RealType(value=200)
+
+        try:
+            # Now test parsing works using the factory name
+            result = cli_api.parse_cli_args(func, ["param=real_type_factory"])
+            assert isinstance(result.param, run.Config)
+            assert result.param.value == 100
+
+            result = cli_api.parse_cli_args(func, ["param=other_factory"])
+            assert isinstance(result.param, run.Config)
+            assert result.param.value == 200
+
+        finally:
+            # Clean up - remove the factory from registry
+            if hasattr(sys.modules[__name__], "real_type_factory"):
+                delattr(sys.modules[__name__], "real_type_factory")
+
 
 class TestListEntrypoints:
     @dataclass

test/cli/test_cli_parser.py

@@ -15,7 +15,16 @@
 
 import sys
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Optional, Type, Union
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Type,
+    Union,
+    ForwardRef,
+)
 
 import pytest
 
@@ -156,6 +165,22 @@ 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):
+        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):
@@ -351,7 +376,8 @@ def func(a: int):
             pass
 
         with pytest.raises(
-            ArgumentValueError, match="Invalid argument: No parameter named 'b' exists for"
+            ArgumentValueError,
+            match="Invalid argument: No parameter named 'b' exists for",
         ):
             parse_cli_args(func, ["b=5"])
 
@@ -401,7 +427,8 @@ def test_parse_int(self):
         assert parse_value("0", int) == 0
         assert parse_value("+789", int) == 789
         with pytest.raises(
-            ParseError, match="Failed to parse '3.14' as <class 'int'>: Invalid integer literal"
+            ParseError,
+            match="Failed to parse '3.14' as <class 'int'>: Invalid integer literal",
         ):
             parse_value("3.14", int)
         with pytest.raises(
@@ -450,7 +477,8 @@ def test_parse_bool(self):
         ):
             parse_value("not_a_bool", bool)
         with pytest.raises(
-            ParseError, match="Failed to parse '2' as <class 'bool'>: Cannot convert .* to bool"
+            ParseError,
+            match="Failed to parse '2' as <class 'bool'>: Cannot convert .* to bool",
         ):
             parse_value("2", bool)
 
@@ -468,7 +496,10 @@ def test_parse_list(self):
 
     def test_parse_dict(self):
         assert parse_value('{"a": 1, "b": 2}', Dict[str, int]) == {"a": 1, "b": 2}
-        assert parse_value('{"x": "foo", "y": "bar"}', Dict[str, str]) == {"x": "foo", "y": "bar"}
+        assert parse_value('{"x": "foo", "y": "bar"}', Dict[str, str]) == {
+            "x": "foo",
+            "y": "bar",
+        }
         assert parse_value("{}", Dict[str, Any]) == {}
         with pytest.raises(ParseError, match="Failed to parse 'not_a_dict' as typing.Dict"):
             parse_value("not_a_dict", Dict[str, int])
@@ -562,7 +593,8 @@ class CustomType:
             pass
 
         with pytest.raises(
-            ParseError, match="Failed to parse 'value' as <class '.*CustomType'>: Unsupported type"
+            ParseError,
+            match="Failed to parse 'value' as <class '.*CustomType'>: Unsupported type",
         ):
             strict_parser.parse("value", CustomType)
 
@@ -602,7 +634,11 @@ def test_parse_constructor(self, parser):
 
     def test_parse_comprehension(self, parser):
         assert parser.parse_comprehension("[x for x in range(3)]") == [0, 1, 2]
-        assert parser.parse_comprehension("{x: x**2 for x in range(3)}") == {0: 0, 1: 1, 2: 4}
+        assert parser.parse_comprehension("{x: x**2 for x in range(3)}") == {
+            0: 0,
+            1: 1,
+            2: 4,
+        }
 
     def test_parse_lambda(self, parser):
         # Test safe lambdas

test/dummy_factory.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass
-from typing import List
+from typing import List, ForwardRef
 
 import nemo_run as run
 
@@ -99,6 +99,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): ...
 
 

test/dummy_type.py

@@ -0,0 +1,3 @@
+class RealType:
+    def __init__(self, value=42):
+        self.value = value