working on handling connections with nested workflows

Author: mahdieh-dst

nipype2pydra/workflow.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import json
 import tempfile
 from pathlib import Path
@@ -10,40 +11,41 @@
 
 
 class WorkflowConverter:
-#creating the wf
+    # creating the wf
     def __init__(self, spec):
         self.spec = spec
-        self.wf = load_class_or_func(self.spec['function'])(
-            **self._parse_workflow_args(self.spec['args'])
-        ) #loads the 'function' in smriprep.yaml, and implement the args (creates a dictionary)
 
-    def node_connections(self, workflow):
+        self.wf = load_class_or_func(self.spec["function"])(
+            **self._parse_workflow_args(self.spec["args"])
+        )  # loads the 'function' in smriprep.yaml, and implement the args (creates a dictionary)
+
+    def node_connections(self, workflow, functions: dict[str, dict], wf_inputs: dict[str, str], wf_outputs: dict[str, str]):
         connections = defaultdict(dict)
 
         # iterates over wf graph, Get connections from workflow graph, store connections in a dictionary
         for edge, props in workflow._graph.edges.items():
             src_node = edge[0].name
             dest_node = edge[1].name
-            for node_conn in props['connect']:
-                src_field = node_conn[1]
-                dest_field = node_conn[0]
-                if src_field.startswith('def'):
-                    print(f"Not sure how to deal with {src_field} in {src_node} to "
-                          f"{dest_node}.{dest_field}")
-                    continue
+            dest_node_fullname = workflow.get_node(dest_node).fullname
+            for node_conn in props["connect"]:
+                src_field = node_conn[0]
+                dest_field = node_conn[1]
+                if src_field.startswith("def"):
+                    functions[dest_node_fullname][dest_field] = src_field
                 else:
-                    src_field = src_field.split('.')[-1]
-                connections[dest_node][dest_field] = f"{src_node}.lzout.{src_field}"
+                    connections[dest_node_fullname][
+                        dest_field
+                    ] = f"{src_node}.lzout.{src_field}"
 
-        # Look for connections in nested workflows via recursion
-        for node in workflow._get_all_nodes():  #TODO: find the method that iterates through the nodes of the workflow (but not nested nodes)---placeholder: find_name_of_appropriate_method
-            if isinstance(node, Workflow):  # TODO: find a way to check whether the node is a standard node or a nested workflow
-                connections.update(self.node_connections(node))
+        for nested_wf in workflow._nested_workflows_cache:
+            connections.update(self.node_connections(nested_wf, functions=functions))
         return connections
 
+
     def generate(self, format_with_black=False):
 
-        connections = self.node_connections(self.wf)
+        functions = defaultdict(dict)
+        connections = self.node_connections(self.wf, functions=functions)
         out_text = ""
         for node_name in self.wf.list_node_names():
             node = self.wf.get_node(node_name)
@@ -59,11 +61,11 @@ def generate(self, format_with_black=False):
                         val = json.dumps(val)
                     except TypeError:
                         pass
-                    if isinstance(val, str) and '\n' in val:
+                    if isinstance(val, str) and "\n" in val:
                         val = '"""' + val + '""""'
                     node_args += f",\n        {arg}={val}"
 
-            for arg, val in connections[node.name].items():
+            for arg, val in connections[node.fullname].items():
                 node_args += f",\n        {arg}=wf.{val}"
 
             out_text += f"""
@@ -72,16 +74,18 @@ def generate(self, format_with_black=False):
 )"""
 
         if format_with_black:
-            out_text = black.format_file_contents(out_text, fast=False, mode=black.FileMode())
+            out_text = black.format_file_contents(
+                out_text, fast=False, mode=black.FileMode()
+            )
         return out_text
 
     @classmethod
     def _parse_workflow_args(cls, args):
         dct = {}
         for name, val in args.items():
-            if isinstance(val, dict) and sorted(val.keys()) == ['args', 'type']:
-                val = load_class_or_func(val['type'])(
-                    **cls._parse_workflow_args(val['args'])
+            if isinstance(val, dict) and sorted(val.keys()) == ["args", "type"]:
+                val = load_class_or_func(val["type"])(
+                    **cls._parse_workflow_args(val["args"])
                 )
             dct[name] = val
         return dct
@@ -92,6 +96,7 @@ def save_graph(self, out_path: Path, format: str = "svg", work_dir: Path = None)
         work_dir = Path(work_dir)
         graph_dot_path = work_dir / "wf-graph.dot"
         self.wf.write_hierarchical_dotfile(graph_dot_path)
-        dot_path = sp.check_output("which dot", shell=True).decode('utf-8').strip()
-        sp.check_call(f"{dot_path} -T{format} {str(graph_dot_path)} > {str(out_path)}",
-                      shell=True)
+        dot_path = sp.check_output("which dot", shell=True).decode("utf-8").strip()
+        sp.check_call(
+            f"{dot_path} -T{format} {str(graph_dot_path)} > {str(out_path)}", shell=True
+        )

tests/test_smriprep.py

@@ -16,7 +16,7 @@ def test_smriprep_conversion(pkg_dir, cli_runner):
         workflow,
         [
             f"{pkg_dir}/example-specs/smriprep.yaml",
-            f"{pkg_dir}/outputs/smriprep_new.py"
+            f"{pkg_dir}/outputs/smriprep_new1.py"
         ]
     )