reworked nested workflow input/output detection and top-level export …

…completes successfully
nipype · tclose · May 17, 2024 · Mar 26, 2024 · Mar 27, 2024 · Mar 28, 2024
commit c32a2b83ad95873e9f374da74d7f5b99636a526c
diff --git a/nipype2pydra/interface/base.py b/nipype2pydra/interface/base.py
@@ -23,14 +23,14 @@
     UsedSymbols,
     types_converter,
     from_dict_converter,
+    unwrap_nested_type,
 )
 from ..statements import (
     ImportStatement,
     parse_imports,
     ExplicitImport,
     from_list_to_imports,
 )
-from fileformats.core.mixin import WithClassifiers
 from fileformats.generic import File
 import nipype2pydra.package
 
@@ -755,14 +755,6 @@ def construct_imports(
                             continue
                     stmts.append(stmt)
 
-        def unwrap_nested_type(t: type) -> ty.List[type]:
-            if issubclass(t, WithClassifiers) and t.is_classified:
-                unwrapped = [t.unclassified]
-                for c in t.classifiers:
-                    unwrapped.extend(unwrap_nested_type(c))
-                return unwrapped
-            return [t]
-
         for tp in itertools.chain(*(unwrap_nested_type(t) for t in nonstd_types)):
             stmts.append(ImportStatement.from_object(tp))
         if include_task:

diff --git a/nipype2pydra/package.py b/nipype2pydra/package.py
@@ -803,11 +803,15 @@ def write_to_module(
         converted_code: ty.Optional[str] = None,
         find_replace: ty.Optional[ty.List[ty.Tuple[str, str]]] = None,
         inline_intra_pkg: bool = False,
+        additional_imports: ty.Optional[ty.List[ImportStatement]] = None,
     ):
         """Writes the given imports, constants, classes, and functions to the file at the given path,
         merging with existing code if it exists"""
         from .helpers import FunctionConverter, ClassConverter
 
+        if additional_imports is None:
+            additional_imports = []
+
         if find_replace is None:
             find_replace = self.find_replace
         else:
@@ -936,6 +940,7 @@ def write_to_module(
             + converter_imports
             + [i for i in used.imports if not i.indent]
             + GENERIC_PYDRA_IMPORTS
+            + additional_imports
         )
 
         if module_fspath.name != "__init__.py":

diff --git a/nipype2pydra/statements/workflow_build.py b/nipype2pydra/statements/workflow_build.py
@@ -17,7 +17,9 @@ class AssignmentStatement:
     indent: str = attrs.field()
     assignments: ty.Dict[str, str] = attrs.field()
 
-    matches_re = re.compile(r"(\s*)(\w[\w\s,]*)\s*=\s*(.*)")
+    matches_re = re.compile(
+        r"(\s*)(\w[\w\s,]*)\s*=\s*(.*)$", flags=re.MULTILINE | re.DOTALL
+    )
 
     def __iter__(self):
         return iter(self.assignments)
@@ -62,7 +64,7 @@ def parse(cls, statement: str) -> "AssignmentStatement":
         )
 
     def __str__(self):
-        return f"{self.indent}{', '.join(self.varnames)} = {', '.join(self.value)}"
+        return f"{self.indent}{', '.join(self.keys())} = {', '.join(self.values())}"
 
 
 @attrs.define
@@ -171,11 +173,17 @@ def targets(self):
 
     @property
     def wf_in(self):
-        return self.source_name is None
+        return self.source_name is None or (
+            (self.target_name, str(self.target_in))
+            in self.workflow_converter._input_mapping
+        )
 
     @property
     def wf_out(self):
-        return self.target_name is None
+        return self.target_name is None or (
+            (self.source_name, str(self.source_out))
+            in self.workflow_converter._output_mapping
+        )
 
     @cached_property
     def conditional(self):
@@ -200,28 +208,28 @@ def wf_in_name(self):
             raise ValueError(
                 f"Cannot get wf_in_name for {self} as it is not a workflow input"
             )
-        source_out_name = (
-            self.source_out
-            if not isinstance(self.source_out, DynamicField)
-            else self.source_out.varname
-        )
-        return self.workflow_converter.input_name(self.source_name, source_out_name)
+        # source_out_name = (
+        #     self.source_out
+        #     if not isinstance(self.source_out, DynamicField)
+        #     else self.source_out.varname
+        # )
+        return self.workflow_converter.get_input(self.source_out, self.source_name).name
 
     @property
     def wf_out_name(self):
         if not self.wf_out:
             raise ValueError(
                 f"Cannot get wf_out_name for {self} as it is not a workflow output"
             )
-        return self.workflow_converter.output_name(self.target_name, self.target_in)
+        return self.workflow_converter.get_output(self.target_in, self.target_name).name
 
     def __str__(self):
         if not self.include:
-            return ""
+            return f"{self.indent}pass\n" if self.conditional else ""
         code_str = ""
         # Get source lazy-field
         if self.wf_in:
-            src = f"{self.workflow_variable}.lzin.{self.wf_in_name}"
+            src = f"{self.workflow_variable}.lzin.{self.source_out}"
         else:
             src = f"{self.workflow_variable}.{self.source_name}.lzout.{self.source_out}"
         if isinstance(self.source_out, DynamicField):
@@ -257,7 +265,7 @@ def __str__(self):
                     f'{intf_name}({self.wf_in_name}={src}, name="{intf_name}"))\n\n'
                 )
                 src = f"{self.workflow_variable}.{intf_name}.lzout.out"
-            code_str += f"{self.indent}{self.workflow_variable}.set_output([({self.wf_out_name!r}, {src})])"
+            code_str += f"{self.indent}{self.workflow_variable}.set_output([({self.target_in!r}, {src})])"
         elif isinstance(self.target_in, VarField):
             code_str += f"{self.indent}setattr({self.workflow_variable}.{self.target_name}.inputs, {self.target_in}, {src})"
         else:
@@ -417,15 +425,15 @@ def converted_interface(self):
 
     def __str__(self):
         if not self.include:
-            return ""
+            return f"{self.indent}pass\n" if self.conditional else ""
         args = ["=".join(a) for a in self.arg_name_vals]
         conn_args = []
         for conn in sorted(self.in_conns, key=attrgetter("target_in")):
             if not conn.include or not conn.lzouttable:
                 continue
             if conn.wf_in:
                 arg = (
-                    f"{conn.target_in}={self.workflow_variable}.lzin.{conn.wf_in_name}"
+                    f"{conn.target_in}={self.workflow_variable}.lzin.{conn.source_out}"
                 )
             else:
                 arg = (
@@ -547,7 +555,7 @@ class AddNestedWorkflowStatement(AddNodeStatement):
 
     def __str__(self):
         if not self.include:
-            return ""
+            return f"{self.indent}pass\n" if self.conditional else ""
         if self.nested_workflow:
             config_params = [
                 f"{n}_{c}={n}_{c}" for n, c in self.nested_workflow.used_configs
@@ -618,26 +626,31 @@ def add_input_connection(self, conn: ConnectionStatement):
             target_name, target_in = conn.target_in.match_to_workflow(
                 self.nested_workflow
             )
-        else:
+        elif isinstance(conn.target_in, NestedVarField):
             target_name = conn.target_in.node_name
             target_in = conn.target_in.varname
+        else:
+            target_in = conn.target_in
+            target_name = None
         if target_name == self.nested_workflow.input_node:
             target_name = None
         nested_input = self.nested_workflow.get_input(target_in, node_name=target_name)
         conn.target_in = nested_input.name
         super().add_input_connection(conn)
         if target_name:
+            # If not connected to the input node, add connections from the nested
+            # workflow input to the target node
             for node in self.nested_workflow.nodes[target_name]:
-                node.add_input_connection(
-                    ConnectionStatement(
-                        source_name=None,
-                        source_out=nested_input.name,
-                        target_name=target_name,
-                        target_in=target_in,
-                        indent=conn.indent,
-                        workflow_converter=self.nested_workflow,
-                    )
+                node_conn = ConnectionStatement(
+                    source_name=None,
+                    source_out=nested_input.name,
+                    target_name=target_name,
+                    target_in=target_in,
+                    indent=conn.indent,
+                    workflow_converter=self.nested_workflow,
                 )
+                self.nested_workflow.connections.append(node_conn)
+                node.add_input_connection(node_conn)
 
     def add_output_connection(self, conn: ConnectionStatement):
         """Adds and output connection to a node, setting as an output of the whole
@@ -658,9 +671,12 @@ def add_output_connection(self, conn: ConnectionStatement):
             source_name, source_out = conn.source_out.match_to_workflow(
                 self.nested_workflow
             )
-        else:
+        elif isinstance(conn.source_out, NestedVarField):
             source_name = conn.source_out.node_name
             source_out = conn.source_out.varname
+        else:
+            source_out = conn.source_out
+            source_name = None
         if source_name == self.nested_workflow.output_node:
             source_name = None
         nested_output = self.nested_workflow.get_output(
@@ -669,17 +685,19 @@ def add_output_connection(self, conn: ConnectionStatement):
         conn.source_out = nested_output.name
         super().add_output_connection(conn)
         if source_name:
+            # If not the output node, add connections to the nested workflow output
+            # from the source node
             for node in self.nested_workflow.nodes[source_name]:
-                node.add_output_connection(
-                    ConnectionStatement(
-                        source_name=source_name,
-                        source_out=source_out,
-                        target_name=None,
-                        target_in=nested_output.name,
-                        indent=conn.indent,
-                        workflow_converter=self.nested_workflow,
-                    )
+                node_conn = ConnectionStatement(
+                    source_name=source_name,
+                    source_out=source_out,
+                    target_name=None,
+                    target_in=nested_output.name,
+                    indent=conn.indent,
+                    workflow_converter=self.nested_workflow,
                 )
+                self.nested_workflow.connections.append(node_conn)
+                node.add_output_connection(node_conn)
 
 
 @attrs.define
@@ -698,11 +716,11 @@ def __str__(self):
         node_name = node.name
         workflow_variable = self.nodes[0].workflow_variable
         if self.is_workflow:
-            nested_wf = node.nested_spec
+            nested_wf = node.nested_workflow
             parts = self.attribute.split(".")
             nested_node_name = parts[2]
             attribute_name = parts[3]
-            target_in = nested_wf.input_name(nested_node_name, attribute_name)
+            target_in = nested_wf.get_input(attribute_name, nested_node_name).name
             attribute = ".".join(parts[:2] + [target_in] + parts[4:])
             workflow_variable = self.nodes[0].workflow_variable
             assert (n.workflow_variable == workflow_variable for n in self.nodes)
@@ -781,14 +799,14 @@ def __str__(self):
             f"    {self.varname} = Workflow("
             f"name={self.workflow_name}, input_spec={{"
             + ", ".join(
-                f"'{i.name}': {i.type}"
+                f"'{i.name}': {i.type.__name__}"
                 for i in sorted(
                     self.workflow_converter.inputs.values(), key=attrgetter("name")
                 )
             )
             + "}, output_spec={"
             + ", ".join(
-                f"'{o.name}': {o.type}"
+                f"'{o.name}': {o.type.__name__}"
                 for o in sorted(
                     self.workflow_converter.outputs.values(), key=attrgetter("name")
                 )
@@ -830,4 +848,6 @@ def __str__(self):
 
     @classmethod
     def parse(cls, statement: str) -> "OtherStatement":
-        return OtherStatement(*re.match(r"(\s*)(.*)", statement).groups())
+        return OtherStatement(
+            *re.match(r"(\s*)(.*)$", statement, flags=re.MULTILINE | re.DOTALL).groups()
+        )
diff --git a/nipype2pydra/utils/__init__.py b/nipype2pydra/utils/__init__.py
@@ -18,6 +18,7 @@
     from_named_dicts_converter,
     str_to_type,
     types_converter,
+    unwrap_nested_type,
     INBUILT_NIPYPE_TRAIT_NAMES,
 )
 from .symbols import (  # noqa: F401

diff --git a/nipype2pydra/utils/misc.py b/nipype2pydra/utils/misc.py
@@ -9,6 +9,7 @@
 from contextlib import contextmanager
 from pathlib import Path
 from fileformats.core import FileSet, from_mime
+from fileformats.core.mixin import WithClassifiers
 from ..exceptions import (
     UnmatchedParensException,
     UnmatchedQuoteException,
@@ -518,3 +519,12 @@ def types_converter(types: ty.Dict[str, ty.Union[str, type]]) -> ty.Dict[str, ty
             tp = str_to_type(tp_or_str)
         converted[name] = tp
     return converted
+
+
+def unwrap_nested_type(t: type) -> ty.List[type]:
+    if issubclass(t, WithClassifiers) and t.is_classified:
+        unwrapped = [t.unclassified]
+        for c in t.classifiers:
+            unwrapped.extend(unwrap_nested_type(c))
+        return unwrapped
+    return [t]