Roll back get_underlying_types change + introduce extract_constituent… (#223)

marcromeyn · web-flow · commit 67de154e4c63 · 2025-05-08T16:46:48.000+02:00
* Roll back get_underlying_types change + introduce extract_constituent_types

Signed-off-by: Marc Romeyn &lt;mromeijn@nvidia.com&gt;

* Bring back test

Signed-off-by: Marc Romeyn &lt;mromeijn@nvidia.com&gt;

---------

Signed-off-by: Marc Romeyn &lt;mromeijn@nvidia.com&gt;
diff --git a/nemo_run/cli/api.py b/nemo_run/cli/api.py
@@ -20,18 +20,22 @@
 import sys
 from dataclasses import dataclass, field
 from functools import cache, wraps
+import typing
 from typing import (
     Any,
+    Annotated,
     Callable,
     Dict,
     Generic,
     List,
     Literal,
     Optional,
     Protocol,
+    Set,
     Tuple,
     Type,
     TypeVar,
+    Union,
     get_args,
     get_type_hints,
     overload,
@@ -62,7 +66,7 @@
     Partial,
     get_nemorun_home,
     get_type_namespace,
-    get_underlying_types,
+    RECURSIVE_TYPES,
 )
 from nemo_run.core.execution import LocalExecutor, SkypilotExecutor, SlurmExecutor
 from nemo_run.core.execution.base import Executor
@@ -475,7 +479,7 @@ def resolve_factory(
     if isinstance(target, str):
         fn = catalogue._get((target, name))
     else:
-        types = get_underlying_types(target)
+        types = extract_constituent_types(target)
         num_missing = 0
         for t in types:
             _namespace = get_type_namespace(t)
@@ -1827,6 +1831,132 @@ def _export_config(output_path: str, format: Optional[str] = None) -> None:
         _export_config(to_json, format="json")
 
 
+def extract_constituent_types(type_hint: Any) -> Set[Type]:
+    """
+    Extract all constituent types from a type hint, including generics and their type arguments.
+    This function recursively traverses complex type hints to find all underlying types.
+
+    For example:
+    - For Union[int, str] -> {int, str}
+    - For List[int] -> {list, int}
+    - For Dict[str, Optional[int]] -> {dict, str, int}
+    - For Callable[..., int] -> {Callable}
+    - For TypeVar('T') -> {TypeVar}
+    - For Annotated[List[int], "metadata"] -> {list, int}
+    - For Optional[List[int]] -> {list, int}
+    - For a class MyClass -> {MyClass}
+
+    The function handles:
+    - Basic types (int, str, etc.)
+    - Generic types (List, Dict, etc.)
+    - Union and Optional types
+    - Annotated types (extracts the underlying type)
+    - TypeVars and ForwardRefs
+    - Callable types
+    - Custom classes
+
+    Args:
+        type_hint: A type hint to analyze. Can be any valid Python type annotation,
+            including complex nested types.
+
+    Returns:
+        A set of all constituent types found in the type hint. NoneType is excluded
+        from the results.
+
+    Note:
+        This function is particularly useful for type checking and reflection,
+        where you need to know all the possible types that could be involved
+        in a type annotation.
+    """
+    # Special case for functions and classes - return the type itself
+    if inspect.isfunction(type_hint) or inspect.isclass(type_hint):
+        return {type_hint}
+
+    # Handle older style type hints (_GenericAlias)
+    if hasattr(typing, "_GenericAlias") and isinstance(type_hint, typing._GenericAlias):  # type: ignore
+        # Correctly handle Annotated by getting the first argument (the actual type)
+        if str(type_hint).startswith("typing.Annotated") or str(type_hint).startswith(
+            "typing_extensions.Annotated"
+        ):
+            # Recurse on the actual type, skipping metadata
+            return extract_constituent_types(type_hint.__args__[0])
+        else:
+            origin = type_hint.__origin__
+
+        if origin in RECURSIVE_TYPES:
+            types = set()
+            for arg in type_hint.__args__:
+                # Add check to skip NoneType here as well
+                if arg is not type(None):
+                    types.update(extract_constituent_types(arg))
+            return types
+        # If not a recursive type handled above, treat it like a concrete generic
+        # Collect types from arguments
+        result = set()
+        for arg in type_hint.__args__:
+            if arg is not type(
+                None
+            ):  # Also skip NoneType here for generics like list[Optional[int]]
+                result.update(extract_constituent_types(arg))
+        # Add the origin itself (e.g., list, dict)
+        if isinstance(origin, type):
+            result.add(origin)
+        # Add the original type_hint if it's a specific generic instantiation (and not a Union/Optional)
+        if origin is not None and origin not in RECURSIVE_TYPES:
+            result.add(type_hint)  # type_hint is the _GenericAlias itself
+        return result  # Return collected types
+
+    # Handle Python 3.9+ style type hints
+    origin = typing.get_origin(type_hint)
+    args = typing.get_args(type_hint)
+
+    # Base case: no origin or args means it's a simple type
+    if origin is None:
+        if type_hint is type(None):
+            return set()
+        if isinstance(type_hint, type):
+            return {type_hint}
+        return {type_hint}  # Return the hint itself if not a type (e.g., TypeVar)
+
+    # Handle Annotated for Python 3.9+
+    if origin is Annotated:
+        # Recurse on the actual type argument, skipping metadata
+        return extract_constituent_types(args[0])
+
+    # Union type (including Optional)
+    if origin is Union:
+        result = set()
+        for arg in args:
+            if arg is not type(None):  # Skip NoneType in Unions
+                result.update(extract_constituent_types(arg))
+        return result
+
+    # List, Dict, etc. - collect types from arguments
+    result = set()
+    for arg in args:
+        result.update(extract_constituent_types(arg))
+
+    # Include the origin type itself if it's a class
+    # This handles both typing module types and Python 3.9+ built-in generic types
+    if isinstance(origin, type):
+        result.add(origin)
+
+    # Add the original type_hint if it's a specific generic instantiation (and not a Union/Annotated)
+    if origin not in (None, Union, Annotated):
+        # type_hint is the original parameterized generic, e.g., List[int]
+        # Add it only if it's indeed a generic (origin of type_hint itself is not None)
+        if typing.get_origin(type_hint) is not None:
+            result.add(type_hint)
+
+    # If no types were added, return the original type hint to preserve behavior
+    if (
+        not result
+    ):  # This covers cases like type_hint being a TypeVar that resulted in an empty set initially
+        return {type_hint}
+
+    return result
+
+
 if __name__ == "__main__":
     app = create_cli()
     app()
diff --git a/nemo_run/cli/cli_parser.py b/nemo_run/cli/cli_parser.py
@@ -1255,7 +1255,8 @@ def parse_factory(parent: Type, arg_name: str, arg_type: Type, value: str) -> An
     """
     import catalogue
 
-    from nemo_run.config import Partial, get_type_namespace, get_underlying_types
+    from nemo_run.config import Partial, get_type_namespace
+    from nemo_run.cli.api import extract_constituent_types
 
     def _get_from_registry(val, annotation, name):
         if catalogue.check_exists(get_type_namespace(annotation), val):
@@ -1297,7 +1298,7 @@ def parse_single_factory(factory_str):
             try:
                 factory_fn = _get_from_registry(factory_name, parent, name=arg_name)
             except catalogue.RegistryError:
-                types = get_underlying_types(arg_type, include_self=True)
+                types = extract_constituent_types(arg_type)
                 for t in types:
                     try:
                         factory_fn = _get_from_registry(factory_name, t, name=factory_name)
diff --git a/nemo_run/config.py b/nemo_run/config.py
@@ -24,7 +24,7 @@
 import typing
 from pathlib import Path
 from types import MappingProxyType
-from typing import Any, Callable, Generic, Iterable, Optional, Type, TypeVar, Union, Set, get_args
+from typing import Any, Callable, Generic, Iterable, Optional, Type, TypeVar, Union, get_args
 
 import fiddle as fdl
 import fiddle._src.experimental.dataclasses as fdl_dc
@@ -103,109 +103,20 @@ def get_type_namespace(typ: Type | Callable) -> str:
     return f"{module}.{_name}"
 
 
-def get_underlying_types(type_hint: Any, include_self: bool = False) -> Set[Type]:
-    """
-    Retrieve the underlying types from a type hint, handling generic types.
-
-    Args:
-        type_hint: The type hint to analyze
-        include_self: If True, include the type_hint itself in the result if it's a specific generic.
-
-    Returns:
-        A set of all underlying types
-    """
-    # Special case for functions and classes - return the type itself
-    if inspect.isfunction(type_hint) or inspect.isclass(type_hint):
-        return {type_hint}
-
-    # Handle older style type hints (_GenericAlias)
-    if hasattr(typing, "_GenericAlias") and isinstance(type_hint, typing._GenericAlias):  # type: ignore
-        # Correctly handle Annotated by getting the first argument (the actual type)
-        if str(type_hint).startswith("typing.Annotated") or str(type_hint).startswith(
-            "typing_extensions.Annotated"
-        ):
-            # Recurse on the actual type, skipping metadata
-            return get_underlying_types(type_hint.__args__[0], include_self=include_self)
+def get_underlying_types(type_hint: typing.Any) -> typing.Set[typing.Type]:
+    if isinstance(type_hint, typing._GenericAlias):  # type: ignore
+        if str(type_hint).startswith("typing.Annotated"):
+            origin = type_hint.__origin__
+            if hasattr(origin, "__origin__"):
+                origin = origin.__origin__
         else:
             origin = type_hint.__origin__
-
         if origin in RECURSIVE_TYPES:
             types = set()
             for arg in type_hint.__args__:
-                # Add check to skip NoneType here as well
-                if arg is not type(None):
-                    types.update(get_underlying_types(arg, include_self=include_self))
+                types.update(get_underlying_types(arg))
             return types
-        # If not a recursive type handled above, treat it like a concrete generic
-        # Collect types from arguments
-        result = set()
-        for arg in type_hint.__args__:
-            if arg is not type(
-                None
-            ):  # Also skip NoneType here for generics like list[Optional[int]]
-                result.update(get_underlying_types(arg, include_self=include_self))
-        # Add the origin itself (e.g., list, dict)
-        if isinstance(origin, type):
-            result.add(origin)
-        # Add the original type_hint if it's a specific generic instantiation (and not a Union/Optional)
-        if include_self and origin is not None and origin not in RECURSIVE_TYPES:
-            result.add(type_hint)  # type_hint is the _GenericAlias itself
-        return result  # Return collected types
-
-    # Handle Python 3.9+ style type hints
-    origin = typing.get_origin(type_hint)
-    args = typing.get_args(type_hint)
-
-    # Base case: no origin or args means it's a simple type
-    if origin is None:
-        if type_hint is type(None):
-            return set()
-        if isinstance(type_hint, type):
-            return {type_hint}
-        return {type_hint}  # Return the hint itself if not a type (e.g., TypeVar)
-
-    # Handle Annotated for Python 3.9+
-    if origin is Annotated:
-        # Recurse on the actual type argument, skipping metadata
-        return get_underlying_types(args[0], include_self=include_self)
-
-    # Union type (including Optional)
-    if origin is typing.Union:
-        result = set()
-        for arg in args:
-            if arg is not type(None):  # Skip NoneType in Unions
-                result.update(get_underlying_types(arg, include_self=include_self))
-        return result
-
-    # List, Dict, etc. - collect types from arguments
-    result = set()
-    for arg in args:
-        result.update(get_underlying_types(arg, include_self=include_self))
-
-    # Include the origin type itself if it's a class
-    # This handles both typing module types and Python 3.9+ built-in generic types
-    if isinstance(origin, type):
-        result.add(origin)
-
-    # Add the original type_hint if it's a specific generic instantiation (and not a Union/Annotated)
-    if (
-        include_self
-        and origin is not None
-        and origin is not typing.Union
-        and origin is not Annotated
-    ):
-        # type_hint is the original parameterized generic, e.g., List[int]
-        # Add it only if it's indeed a generic (origin of type_hint itself is not None)
-        if typing.get_origin(type_hint) is not None:
-            result.add(type_hint)
-
-    # If no types were added, return the original type hint to preserve behavior
-    if (
-        not result
-    ):  # This covers cases like type_hint being a TypeVar that resulted in an empty set initially
-        return {type_hint}
-
-    return result
+    return {type_hint}
 
 
 def from_dict(raw_data: dict | list | str | float | int | bool, cls: Type[_T]) -> _T:
diff --git a/test/cli/test_api.py b/test/cli/test_api.py
@@ -42,6 +42,7 @@
     _load_workspace_file,
     _load_workspace,
     main as cli_main,
+    extract_constituent_types,
 )
 from test.dummy_factory import DummyModel, dummy_entrypoint
 import nemo_run.cli.cli_parser  # Import the module to mock its function
@@ -1861,3 +1862,32 @@ def test_format_help_with_docs_flag(
         with patch("sys.argv", ["script.py", "test_cmd", "--help", "-d"]):
             cmd.format_help(mock_ctx, mock_formatter)
         entrypoint.help.assert_called_once_with(mock_console, with_docs=True)
+
+
+class TestExtractConstituentTypes:
+    @pytest.mark.parametrize(
+        "type_hint, expected_types",
+        [
+            (int, {int}),
+            (str, {str}),
+            (bool, {bool}),
+            (float, {float}),
+            (list[int], {list, int, list[int]}),
+            (dict[str, float], {dict, str, float, dict[str, float]}),
+            (Union[int, str], {int, str}),
+            (Optional[int], {int}),  # Optional[T] is Union[T, NoneType]
+            (list[Union[int, str]], {list, int, str, list[Union[int, str]]}),
+            (dict[str, list[int]], {dict, str, list, int, dict[str, list[int]], list[int]}),
+            (Optional[list[str]], {list, str, list[str]}),
+            (Annotated[int, "meta"], {int}),
+            (Annotated[list[str], "meta"], {list, str, list[str]}),
+            (Annotated[Optional[dict[str, bool]], "meta"], {dict, str, bool, dict[str, bool]}),
+            (Union[Annotated[int, "int_meta"], Annotated[str, "str_meta"]], {int, str}),
+            (DummyModel, {DummyModel}),
+            (Optional[DummyModel], {DummyModel}),
+            (list[DummyModel], {list, DummyModel, list[DummyModel]}),
+        ],
+    )
+    def test_various_type_hints(self, type_hint, expected_types):
+        """Test get_underlying_types with various type hints."""
+        assert extract_constituent_types(type_hint) == expected_types
diff --git a/test/run/torchx_backend/schedulers/test_slurm.py b/test/run/torchx_backend/schedulers/test_slurm.py
@@ -297,7 +297,6 @@ def test_tunnel_log_iterator():
 
 
 @mock.patch("nemo_run.run.torchx_backend.schedulers.slurm.SLURM_JOB_DIRS", "mock_job_dirs_path")
-@pytest.mark.xfail
 def test_get_job_dirs():
     # Single test using direct file manipulation instead of complex mocks
     with tempfile.TemporaryDirectory() as temp_dir:
diff --git a/test/test_config.py b/test/test_config.py
