pipelime-python 2.3.0__py3-none-any.whl → 3.0.0__py3-none-any.whl

pipelime/__init__.py

@@ -2,4 +2,4 @@
 
 __author__ = "Eyecan.ai"
 __email__ = "info@eyecan.ai"
-__version__ = "2.3.0"
+__version__ = "3.0.0"

pipelime/choixe/ast/nodes.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
+import functools
+import types
+import typing
 from abc import ABC, abstractmethod
-from dataclasses import field
+from dataclasses import dataclass, field, fields
 from typing import Any, Dict, Optional, Sequence, Tuple, Union
 
-from pydantic.v1.dataclasses import dataclass
-
 
 class NodeVisitor:  # pragma: no cover
     """Base class for all Choixe visitors. A visitor can be seen as a special recursive
@@ -87,10 +88,47 @@ class NodeVisitor:  # pragma: no cover
         return self._ignore(node)
 
 
+@functools.lru_cache(maxsize=None)
+def _node_typed_fields(cls: type) -> Tuple[Tuple[str, Tuple[type, ...], bool], ...]:
+    """For each dataclass field of `cls` whose resolved annotation is a `Node`
+    subclass, `Optional[<Node subclass>]`, or a `Union` of `Node` subclasses,
+    returns `(field_name, allowed_node_types, allow_none)`. Container-typed
+    fields (e.g. `Dict[...]`, `Sequence[...]`) and `Any` are left out. Cached
+    per class (outside any instance, since `HashNode.__hash__` hashes
+    `self.__dict__.values()`)."""
+    hints = typing.get_type_hints(cls)
+    specs = []
+    for f in fields(cls):
+        hint = hints.get(f.name)
+        is_union = typing.get_origin(hint) is Union or isinstance(hint, types.UnionType)
+        candidates = typing.get_args(hint) if is_union else (hint,)
+        node_types = tuple(
+            c for c in candidates if isinstance(c, type) and issubclass(c, Node)
+        )
+        if node_types:
+            specs.append((f.name, node_types, type(None) in candidates))
+    return tuple(specs)
+
+
 @dataclass
 class Node(ABC):
     """A generic element of the Choixe AST, all nodes must implement this interface."""
 
+    def __post_init__(self) -> None:
+        for name, node_types, allow_none in _node_typed_fields(type(self)):
+            value = getattr(self, name, None)
+            if value is None and allow_none:
+                continue
+            if isinstance(value, node_types):
+                continue
+            expected = " or ".join(t.__name__ for t in node_types)
+            if allow_none:
+                expected += " or None"
+            raise TypeError(
+                f"{type(self).__name__}.{name} expected {expected}, "
+                f"got {type(value).__name__}"
+            )
+
     @abstractmethod
     def accept(self, visitor: NodeVisitor) -> Any:
         """Accepts an incoming visitor. This method should simply call the respective

pipelime/choixe/utils/common.py

@@ -1,5 +1,39 @@
+import re
 import typing as t
 
+# pydash's string path delimiter (the same in pydash 8.0 and 8.1): an unescaped `.`,
+# or a `[<int>]` list index (captured, so that it is kept as a path part)
+_PYDASH_PATH_DELIM = re.compile(r"(?<!\\)(?:\\\\)*\.|(\[-?\d+\])")
+
+
+def pydash_path(key_path: t.Any) -> t.Any:
+    r"""A key path for pydash (``get``, ``set_``, ``has``, ``to_path``) addressing the
+    same value on every pydash version.
+
+    pydash 8.1 reads the empty key of a leading, trailing or doubled ``.`` in a string
+    path (``".a"`` is ``["", "a"]``, ``"."`` is ``["", ""]``), where earlier versions
+    dropped it (``["a"]``, ``[]``: the whole object). pipelime keeps the pre-8.1 meaning:
+    the empty keys of a dotted or indexed string path are dropped here, and the other
+    parts, escapes included, are joined back into a canonical path (``"a.b[0].c"``) that
+    every pydash version splits the same way. A path left with a single key is returned
+    unescaped (pydash takes a plain key as is), with no key as ``[]``. Any other path (a
+    plain key without ``.`` or ``[``, a key list, an int) is returned unchanged.
+    """
+    if not (isinstance(key_path, str) and ("." in key_path or "[" in key_path)):
+        return key_path
+    # `re.split` alternates the text between two delimiters (a key, "" if empty) and
+    # the captured group (a `[<int>]` list index, None for a `.`)
+    parts = _PYDASH_PATH_DELIM.split(key_path)
+    path = ""
+    for i, part in enumerate(parts):
+        if part:
+            path += part if i % 2 or not path else "." + part
+    if not path:
+        return []
+    if "." in path or "[" in path:
+        return path
+    return path.replace("\\\\", "\\")  # pydash's unescape; no `\.` is left here
+
 
 class _RefItem:
     def __init__(self, parent, idx):
@@ -73,7 +107,7 @@ def deep_set_(  # noqa: C901
     from pydash.utilities import to_path
 
     if key_path is not None:  # pragma: no branch
-        key_path_tokens = to_path(key_path)
+        key_path_tokens = to_path(pydash_path(key_path))
         parent_node = _RefItem(obj, None)
         for tk in key_path_tokens:
             if isinstance(tk, int):

pipelime/choixe/visitors/decoder.py

@@ -3,7 +3,7 @@ from pathlib import Path
 from typing import Any
 
 import numpy as np
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from pipelime.choixe.ast.nodes import LiteralNode, Node
 from pipelime.choixe.visitors.unparser import Unparser
@@ -20,7 +20,10 @@ class Decoder(Unparser):
             return data.item()
         elif isinstance(data, BaseModel):
             symbol = f"{data.__module__}.{data.__class__.__qualname__}"
-            return {"$model": symbol, "$args": json.loads(data.json(by_alias=True))}
+            return {
+                "$model": symbol,
+                "$args": json.loads(data.model_dump_json(by_alias=True)),
+            }
         elif isinstance(data, Path):
             return str(data)
         else:

pipelime/choixe/visitors/inspector.py

@@ -10,6 +10,7 @@ import pydash as py_
 
 import pipelime.choixe.ast.nodes as ast
 from pipelime.choixe.ast.parser import parse
+from pipelime.choixe.utils.common import pydash_path
 from pipelime.choixe.utils.io import load
 
 
@@ -87,11 +88,11 @@ class Inspector(ast.NodeVisitor):
 
             id_ = node.identifier.data  # type: ignore
             default = None if node.default is None else unparse(node.default)
-            variables = py_.set_({}, id_, default)
+            variables = py_.set_({}, pydash_path(id_), default)
 
             help_strings = {}
             if node.help is not None:
-                help_strings = py_.set_({}, id_, node.help.data)
+                help_strings = py_.set_({}, pydash_path(id_), node.help.data)
 
             environ = {}
             if node.env is not None and node.env.data:  # type: ignore
@@ -150,7 +151,7 @@ class Inspector(ast.NodeVisitor):
         if node.identifier is not None:  # pragma: no branch
             self._named_for_loops[node.identifier.data] = node.iterable.data
         iterable_insp = (
-            Inspection(variables=py_.set_({}, node.iterable.data, None))
+            Inspection(variables=py_.set_({}, pydash_path(node.iterable.data), None))
             if isinstance(node.iterable.data, str)
             else Inspection(processed=True)
         )
@@ -160,7 +161,9 @@ class Inspector(ast.NodeVisitor):
     def visit_switch(self, node: ast.SwitchNode) -> Inspection:
         insp = node.value.accept(self)
         if insp.processed:  # pragma: no branch
-            insp = Inspection(variables=py_.set_({}, node.value.data, None))  # type: ignore
+            insp = Inspection(
+                variables=py_.set_({}, pydash_path(node.value.data), None)  # type: ignore
+            )
         cases_insp = sum(
             [x.accept(self) + y.accept(self) for x, y in node.cases],
             start=Inspection(processed=True),
@@ -186,7 +189,7 @@ class Inspector(ast.NodeVisitor):
                 if key:
                     iterable_name = self._named_for_loops[loop_id]
                     full_path = f"{iterable_name}.{key}"
-                    py_.set_(variables, full_path, None)
+                    py_.set_(variables, pydash_path(full_path), None)
             insp = insp + sub_insp + Inspection(variables=variables)
         return insp
 

pipelime/choixe/visitors/processor.py

@@ -11,6 +11,7 @@ import pydash as py_
 
 import pipelime.choixe.ast.nodes as ast
 from pipelime.choixe.ast.parser import parse
+from pipelime.choixe.utils.common import pydash_path
 from pipelime.choixe.utils.imports import import_symbol
 from pipelime.choixe.utils.io import PipelimeTmp, load
 from pipelime.choixe.utils.rand import rand
@@ -148,8 +149,8 @@ class Processor(ast.NodeVisitor):
             var_value = None
             found = False
 
-            if py_.has(self._context, id_):
-                var_value = py_.get(self._context, id_)
+            if py_.has(self._context, pydash_path(id_)):
+                var_value = py_.get(self._context, pydash_path(id_))
                 found = True
             if not found and env:
                 value = os.getenv(id_)
@@ -214,14 +215,24 @@ class Processor(ast.NodeVisitor):
         symbol_branches = node.symbol.accept(self)
         args_branches = node.args.accept(self)
         branches = self._branches(symbol_branches, args_branches)
-        return [import_symbol(s, cwd=self._cwd).parse_obj(a) for s, a in branches]
+        models = []
+        for s, a in branches:
+            model_cls = import_symbol(s, cwd=self._cwd)
+            if not hasattr(model_cls, "model_validate"):
+                raise TypeError(
+                    f"`$model` needs a pydantic v2 model, got {model_cls!r}: "
+                    "pipelime 3 no longer supports pydantic.v1 models "
+                    "(see docs/migration/pydantic_v2.md)"
+                )
+            models.append(model_cls.model_validate(a))
+        return models
 
     def visit_for(self, node: ast.ForNode) -> List[Any]:
         if isinstance(node.iterable.data, str):
-            if not py_.has(self._context, node.iterable.data):
+            if not py_.has(self._context, pydash_path(node.iterable.data)):
                 iterable = self._handle_missing_for(node.iterable.data)
             else:
-                iterable = py_.get(self._context, node.iterable.data)
+                iterable = py_.get(self._context, pydash_path(node.iterable.data))
         else:
             iterable = node.iterable.data
 
@@ -278,10 +289,10 @@ class Processor(ast.NodeVisitor):
         for branch in all_branches:
             varname = branch[0]
 
-            if not py_.has(self._context, varname):
+            if not py_.has(self._context, pydash_path(varname)):
                 value = self._handle_missing_switch(varname)
             else:
-                value = py_.get(self._context, varname)
+                value = py_.get(self._context, pydash_path(varname))
 
             # Match the value to the correct case
             for i in range(len(node.cases)):
@@ -321,7 +332,7 @@ class Processor(ast.NodeVisitor):
         items = []
         for branch in branches:
             loop_id, _, key = str(branch).partition(sep)  # type: ignore
-            item = py_.get(self._loop_data[loop_id].item, f"{sep}{key}")
+            item = py_.get(self._loop_data[loop_id].item, pydash_path(f"{sep}{key}"))
             items.append(item)
         return items
 

pipelime/choixe/xconfig.py

@@ -9,7 +9,7 @@ from schema import Schema
 
 from pipelime.choixe.ast.nodes import Node
 from pipelime.choixe.ast.parser import parse
-from pipelime.choixe.utils.common import deep_set_
+from pipelime.choixe.utils.common import deep_set_, pydash_path
 from pipelime.choixe.utils.io import dump, load
 from pipelime.choixe.visitors import Inspection, decode, inspect, process, walk
 
@@ -139,7 +139,7 @@ class XConfig(Box):
         Returns:
             Any: The value at the specified path.
         """
-        return py_.get(self, full_key, default=default)
+        return py_.get(self, pydash_path(full_key), default=default)
 
     def deep_set(
         self,
@@ -159,7 +159,7 @@ class XConfig(Box):
                 Defaults to False.
         """
 
-        if not only_valid_keys or py_.has(self, full_key):
+        if not only_valid_keys or py_.has(self, pydash_path(full_key)):
             deep_set_(self, key_path=full_key, value=value, append=append_values)
 
     def deep_update(

pipelime/cli/main.py

@@ -6,7 +6,6 @@ import typing as t
 from pathlib import Path
 
 import typer
-from pydantic.v1 import BaseModel
 
 from pipelime.choixe import XConfig
 from pipelime.choixe.visitors.processor_ui import ProcessorUi
@@ -18,12 +17,15 @@ from pipelime.cli.utils import (
     print_command_op_stage_info,
     print_commands_ops_stages_list,
 )
+from pipelime.utils.pydantic_compat import PipelimeModel
 
 if t.TYPE_CHECKING:
     from pipelime.piper.checkpoint import Checkpoint
 
 
-class PlCliOptions(BaseModel):
+class PlCliOptions(PipelimeModel):
+    # a `PipelimeModel` for the v1 Optional rule: the options saved in a checkpoint
+    # may lack an `Optional` key
     _namespace: t.ClassVar[str] = "__plmain"
 
     config: t.List[Path]
@@ -43,7 +45,7 @@ class PlCliOptions(BaseModel):
     pipelime_tmp: t.Optional[str]
 
     def purged_dict(self):
-        return self._purge(self.dict())
+        return self._purge(self.model_dump())
 
     def _purge(self, data):
         if isinstance(data, Path):
@@ -918,14 +920,19 @@ def run_command(
 
     import time
 
-    from pydantic.v1.error_wrappers import ValidationError
+    from pydantic import ValidationError
+    from rich.markup import escape
 
     from pipelime.choixe.utils.io import PipelimeTmp, dump
-    from pipelime.cli.pretty_print import print_command_outputs, print_info
+    from pipelime.cli.pretty_print import (
+        print_command_outputs,
+        print_error,
+        print_info,
+    )
     from pipelime.cli.tui import TuiApp, is_tui_needed
     from pipelime.cli.utils import (
+        format_validation_error,
         get_pipelime_command_cls,
-        show_field_alias_valerr,
         time_to_str,
     )
     from pipelime.commands import TempCommand
@@ -972,12 +979,13 @@ def run_command(
             # so let's show the tui again if it was needed in the first place
 
     except ValidationError as e:
-        show_field_alias_valerr(e)
+        # plain text, not Rich markup: `[type=...]` and bracketed values must survive
+        print_error(escape(format_validation_error(e, cmd_cls)))
         raise e
 
     if verbose > 0:
         print_info(f"\nCreated command `{command}`:")
-        print_info(cmd_obj.dict(), pretty=True)
+        print_info(cmd_obj.model_dump(), pretty=True)
 
     if dry_run or verbose > 0:
         print_info(f"\nRunning `{command}`...")

pipelime/cli/pretty_print.py

@@ -1,7 +1,7 @@
 import inspect
 import typing as t
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 from rich import box, get_console
 from rich import print as rprint
 from rich.markup import escape
@@ -9,6 +9,14 @@ from rich.panel import Panel
 from rich.pretty import Pretty
 from rich.table import Column, Table
 
+from pipelime.utils.pydantic_compat import (
+    FieldView,
+    iter_fields,
+    model_title,
+    strip_annotated,
+    type_info,
+)
+
 if t.TYPE_CHECKING:
     from pipelime.cli.utils import ActionInfo
 
@@ -92,9 +100,7 @@ def show_spinning_status(text: str):
 
 
 def get_model_title(model_cls: t.Type[BaseModel]) -> str:
-    if model_cls.__config__.title:
-        return model_cls.__config__.title
-    return model_cls.__name__
+    return model_title(model_cls)
 
 
 def get_model_classpath(model_cls: t.Type[BaseModel]) -> str:
@@ -104,17 +110,15 @@ def get_model_classpath(model_cls: t.Type[BaseModel]) -> str:
 
 
 def print_model_field_values(
-    model_fields: t.Mapping,
+    model_fields: t.Mapping[str, FieldView],
     port_values: t.Mapping[str, t.Any],
     icon: str = "",
 ):
     for k, v in port_values.items():
         rprint(f"\n{icon if icon else '***'} {k}:")
         # Ports might be virtual, as in ShellCommand, so they might not be in the model
-        if k in model_fields and model_fields[k].field_info.description:
-            rprint(
-                f"[italic grey50]{escape(model_fields[k].field_info.description)}[/]"
-            )
+        if k in model_fields and model_fields[k].description:
+            rprint(f"[italic grey50]{escape(model_fields[k].description)}[/]")
         rprint(
             "[green]"
             + escape(str(v) if isinstance(v, (bytes, str)) else repr(v))
@@ -122,12 +126,20 @@ def print_model_field_values(
         )
 
 
+def _fields_by_name(model_cls) -> t.Dict[str, FieldView]:
+    return {f.name: f for f in iter_fields(model_cls)}
+
+
 def print_command_inputs(command: "PipelimeCommand"):  # type: ignore # noqa: E602,F821
-    print_model_field_values(command.__fields__, command.get_inputs(), _input_icon())
+    print_model_field_values(
+        _fields_by_name(type(command)), command.get_inputs(), _input_icon()
+    )
 
 
 def print_command_outputs(command: "PipelimeCommand"):  # type: ignore # noqa: E602,F821
-    print_model_field_values(command.__fields__, command.get_outputs(), _output_icon())
+    print_model_field_values(
+        _fields_by_name(type(command)), command.get_outputs(), _output_icon()
+    )
 
 
 def print_actions_short_help(*actions_info: "ActionInfo", show_class_path: bool = True):
@@ -218,7 +230,7 @@ def print_model_info(
 
 def _field_row(
     grid: Table,
-    field,
+    field: FieldView,
     indent: int,
     indent_offs: int,
     show_piper_port: bool,
@@ -227,35 +239,29 @@ def _field_row(
 ):
     from enum import Enum
 
-    expand_help = field.field_info.extra.get("expand_help", False)
+    expand_help = field.extra.get("expand_help", False)
 
-    is_model = _is_model(field.outer_type_) and not inspect.isabstract(
-        field.outer_type_
-    )
+    inner_type = field.inner_type
+    is_model = field.is_model and not inspect.isabstract(inner_type)
 
     if show_description:
-        # NB: docs should not come from the inner __root__ type
-        if field.field_info.description:
-            field_docs = field.field_info.description
-        elif hasattr(field.outer_type_, "__doc__") and field.outer_type_.__doc__:
-            field_docs = str(inspect.getdoc(field.outer_type_))
+        # NB: docs should not come from the inner root type
+        if field.description:
+            field_docs = field.description
+        elif hasattr(inner_type, "__doc__") and inner_type.__doc__:
+            field_docs = str(inspect.getdoc(inner_type))
         else:
             field_docs = ""
 
-        # field_docs = " ".join(field_docs.split())
-
-    has_root_item = ("__root__" in field.outer_type_.__fields__) if is_model else False
-    field_outer_type = (
-        field.outer_type_.__fields__["__root__"].outer_type_
-        if has_root_item
-        else field.outer_type_
-    )
+    root_type = field.root_type if is_model else None
+    has_root_item = root_type is not None
+    field_outer_type = root_type if has_root_item else inner_type
 
     if show_piper_port:
         from pipelime.piper import PiperPortType
 
         fport = str(
-            field.field_info.extra.get("piper_port", PiperPortType.PARAMETER).value
+            field.extra.get("piper_port", PiperPortType.PARAMETER).value
         ).upper()
 
         if fport == PiperPortType.INPUT.value.upper():
@@ -271,10 +277,10 @@ def _field_row(
         + ("[bold dark_orange]" if indent == 0 else "")
         + (
             f"{escape(field.name)} / "
-            if field.model_config.allow_population_by_field_name and field.has_alias
+            if field.populate_by_name and field.has_alias
             else ""
         )
-        + f"{escape(field.alias)}"
+        + f"{escape(field.effective_alias)}"
         + ("[/]" if indent == 0 else "")
     ]
 
@@ -294,10 +300,13 @@ def _field_row(
         line.append(fport)  # type: ignore
 
     # Default value
-    field_default = field.get_default()
-    if isinstance(field_default, Enum):
-        field_default = field_default.value
-    line.append("[red]✗[/]" if field.required else f"[green]{field_default}[/]")
+    if field.required:
+        line.append("[red]✗[/]")
+    else:
+        field_default = field.default
+        if isinstance(field_default, Enum):
+            field_default = field_default.value
+        line.append(f"[green]{field_default}[/]")
 
     grid.add_row(*line)
 
@@ -381,12 +390,9 @@ def _get_signature(model_cls: t.Type[BaseModel]) -> str:
 
             return f"\n  {formatted}"
 
-    fullname = {mfield.alias: mfield.name for mfield in model_cls.__fields__.values()}
-    excluded = [
-        mfield.alias
-        for mfield in model_cls.__fields__.values()
-        if mfield.field_info.exclude
-    ]
+    fields = list(iter_fields(model_cls))
+    fullname = {f.effective_alias: f.name for f in fields}
+    excluded = [f.effective_alias for f in fields if f.exclude]
 
     sig = inspect.signature(model_cls)
     sig = sig.replace(
@@ -418,48 +424,37 @@ def _is_model(type_):
     return inspect.isclass(type_) and issubclass(type_, BaseModel)
 
 
+_TYPING_MODULES = ("typing", "typing_extensions")
+
+
 def _human_readable_type(field_outer_type):
     from enum import Enum
 
-    from pydantic.v1.typing import (
-        WithArgsTypes,
-        get_args,
-        get_origin,
-        is_union,
-        typing_base,  # noqa: F401  # type: ignore
-    )
+    # constrained types (`Annotated[int, Gt(gt=0)]`) are shown as their base type
+    v = strip_annotated(field_outer_type)
+    ti = type_info(v)
 
-    v = field_outer_type
-    if (
-        not isinstance(v, typing_base)
-        and not isinstance(v, WithArgsTypes)
-        and not isinstance(v, type)
-    ):
-        v = v.__class__
-
-    v_orig = get_origin(v)
-
-    if is_union(v_orig):
-        return " | ".join(map(_human_readable_type, get_args(v)))
-    if inspect.isclass(v_orig):
-        if issubclass(dict, v_orig):
-            return "{" + ": ".join(map(_human_readable_type, get_args(v))) + "}"
-        if issubclass(list, v_orig):
-            return "[" + ", ".join(map(_human_readable_type, get_args(v))) + ", ...]"
-        if issubclass(tuple, v_orig):
-            return "(" + ", ".join(map(_human_readable_type, get_args(v))) + ")"
+    if ti.is_union:
+        return " | ".join(map(_human_readable_type, ti.args))
+    if inspect.isclass(ti.origin):
+        if issubclass(dict, ti.origin):
+            return "{" + ": ".join(map(_human_readable_type, ti.args)) + "}"
+        if issubclass(list, ti.origin):
+            return "[" + ", ".join(map(_human_readable_type, ti.args)) + ", ...]"
+        if issubclass(tuple, ti.origin):
+            return "(" + ", ".join(map(_human_readable_type, ti.args)) + ")"
 
     if inspect.isclass(v) and issubclass(v, Enum):
         v = v.__name__ + "{" + ", ".join(e.name.lower() for e in v) + "}"
-    elif isinstance(v, WithArgsTypes):
-        # Generic alias are constructs like `list[int]`
+    elif ti.origin is not None:
+        # generic aliases (`Sequence[int]`, `Literal[...]`)
         v = str(v).replace("typing.", "")
     else:
-        try:
-            v = v.__name__
-        except AttributeError:
-            # happens with typing objects
-            v = str(v).replace("typing.", "")
+        if not inspect.isclass(v) and type(v).__module__ not in _TYPING_MODULES:
+            # a value among the type args, e.g. `...` in `Tuple[int, ...]`
+            v = type(v)
+        # classes and named typing objects (`TypeVar`) by name, the others as written
+        v = getattr(v, "__name__", None) or str(v).replace("typing.", "")
 
     return v.replace("NoneType", "None")
 
@@ -497,8 +492,8 @@ def _iterate_model_fields(
     add_blank_row,
 ):
     no_data = True
-    for field in model_cls.__fields__.values():  # type: ignore
-        if not field.field_info.exclude:
+    for field in iter_fields(model_cls):
+        if not field.exclude:
             no_data = False
             _field_row(
                 grid,

pipelime/cli/tui/tui.py

@@ -20,6 +20,7 @@ from pipelime.cli.tui.utils import (
 from pipelime.cli.utils import parse_user_input
 from pipelime.piper import PipelimeCommand
 from pipelime.stages import StageInput
+from pipelime.utils.pydantic_compat import iter_fields
 
 
 @dataclass(frozen=True)
@@ -143,8 +144,8 @@ class TuiApp(App[Mapping]):
         """
         tui_fields = {}
 
-        for field in self.cmd_cls.__fields__.values():
-            if field.type_ == StageInput:
+        for field in iter_fields(self.cmd_cls):
+            if field.inner_type is StageInput:
                 tui_fields[field.name] = init_stageinput_tui_field(field, cmd_args)
             else:
                 tui_fields[field.name] = init_tui_field(field, cmd_args)