handling special super methods

Author: tclose

nipype2pydra/interface/base.py

@@ -14,7 +14,7 @@
 import attrs
 from attrs.converters import default_if_none
 import nipype.interfaces.base
-from nipype.interfaces.base import traits_extension, CommandLine
+from nipype.interfaces.base import traits_extension, CommandLine, BaseInterface
 from pydra.engine import specs
 from pydra.engine.helpers import ensure_list
 from ..utils import (
@@ -1019,7 +1019,8 @@ def _misc_cleanups(self, body: str) -> str:
             body,
             flags=re.MULTILINE,
         )
-        body = re.sub(r"\w+runtime\.(stdout|stderr)", r"\1", body)
+        body = re.sub(r"\bruntime\.(stdout|stderr)", r"\1", body)
+        body = re.sub(r"\boutputs\.(\w+)", r"outputs['\1']", body)
         body = body.replace("os.getcwd()", "output_dir")
         return body
 
@@ -1307,22 +1308,30 @@ def replace_super(match):
 
         return re.sub(r"super\([^\)]*\)\.(\w+)\(([^\)]*)\)", replace_super, method_body)
 
-    def unwrap_nested_methods(self, method_body, additional_args=()):
+    def unwrap_nested_methods(
+        self, method_body, additional_args=(), inputs_as_dict: bool = False
+    ):
         """
         Converts nested method calls into function calls
         """
         # Add args to the function signature of method calls
         method_re = re.compile(r"self\.(\w+)(?=\()", flags=re.MULTILINE | re.DOTALL)
         method_names = [m.__name__ for m in self.referenced_methods]
-        unrecognised_methods = set(
+        method_body = strip_comments(method_body)
+        omitted_methods = {}
+        for method_name in set(
             m for m in method_re.findall(method_body) if m not in method_names
-        )
-        assert (
-            not unrecognised_methods
-        ), f"Found the following unrecognised methods {unrecognised_methods}"
+        ):
+            omitted_methods[method_name] = find_super_method(
+                self.nipype_interface, method_name
+            )[0]
         splits = method_re.split(method_body)
         new_body = splits[0]
         for name, args in zip(splits[1::2], splits[2::2]):
+            if name in omitted_methods:
+                new_body += self.SPECIAL_SUPER_MAPPINGS[omitted_methods[name]]
+                new_body += extract_args(args)[-1][1:]
+                continue
             # Assign additional return values (which were previously saved to member
             # attributes) to new variables from the method call
             if self.method_returns[name]:
@@ -1342,18 +1351,18 @@ def unwrap_nested_methods(self, method_body, additional_args=()):
                     else:
                         new_body += ",".join(self.method_returns[name]) + " = "
                 else:
-                    raise NotImplementedError(
+                    logger.warning(
                         "Could not augment the return value of the method converted from "
-                        "a function with the previously assigned attributes as it is used "
-                        "directly. Need to replace the method call with a variable and "
-                        "assign the return value to it on a previous line"
+                        f"a function '{name}' with the previously assigned attributes "
+                        f"{self.method_returns[name]} as the method doesn't have a "
+                        f"singular return statement at the end of the method"
                     )
             # Insert additional arguments to the method call (which were previously
             # accessed via member attributes)
             new_body += name + insert_args_in_signature(
                 args,
                 [
-                    f"{a}={a}"
+                    f"{a}=inputs['{a}']" if inputs_as_dict else f"{a}={a}"
                     for a in (list(self.method_args[name]) + list(additional_args))
                 ],
             )
@@ -1368,6 +1377,7 @@ def unwrap_nested_methods(self, method_body, additional_args=()):
     SPECIAL_SUPER_MAPPINGS = {
         CommandLine._list_outputs: "{}",
         CommandLine._format_arg: "argstr.format(**inputs)",
+        BaseInterface._check_version_requirements: "[]",
     }
 
     INPUT_KEYS = [
@@ -1433,3 +1443,7 @@ def find_super_method(
         f"Could not find super of '{method_name}' method in base classes of "
         f"{super_base}"
     )
+
+
+def strip_comments(src: str) -> str:
+    return re.sub(r"\s*#.*", "", src)

nipype2pydra/interface/shell_command.py

@@ -288,11 +288,10 @@ def format_arg_code(self):
         )
         if not body:
             return ""
-        body = self.unwrap_nested_methods(body)
+        body = self.unwrap_nested_methods(body, inputs_as_dict=True)
         body = self.replace_supers(body)
 
-        code_str = f"""def _format_arg({name_arg}, {val_arg}, inputs, argstr):
-    parsed_inputs = _parse_inputs(inputs) if inputs else {{}}
+        code_str = f"""def _format_arg({name_arg}, {val_arg}, inputs, argstr):{self.parse_inputs_call}
     if {val_arg} is None:
         return ""
 {body}
@@ -324,7 +323,7 @@ def parse_inputs_code(self) -> str:
 
         # Strip out return value
         body = re.sub(r"\s*return .*\n", "", body)
-        body = self.unwrap_nested_methods(body)
+        body = self.unwrap_nested_methods(body, inputs_as_dict=True)
         body = self.replace_supers(body)
 
         code_str = "def _parse_inputs(inputs):\n    parsed_inputs = {{}}"
@@ -352,11 +351,10 @@ def defaults_code(self):
 
         if not body:
             return ""
-        body = self.unwrap_nested_methods(body)
+        body = self.unwrap_nested_methods(body, inputs_as_dict=True)
         body = self.replace_supers(body)
 
-        code_str = f"""def _gen_filename(name, inputs):
-    parsed_inputs = _parse_inputs(inputs) if inputs else {{}}
+        code_str = f"""def _gen_filename(name, inputs):{self.parse_inputs_call}
 {body}
 """
         # Create separate default function for each input field with genfile, which
@@ -376,31 +374,39 @@ def callables_code(self):
 
         if not self.callable_output_fields:
             return ""
-
-        body = _strip_doc_string(
-            inspect.getsource(self.nipype_interface._list_outputs).split("\n", 1)[-1]
-        )
+        if hasattr(self.nipype_interface, "aggregate_outputs"):
+            func_name = "aggregate_outputs"
+            body = _strip_doc_string(
+                inspect.getsource(self.nipype_interface.aggregate_outputs).split(
+                    "\n", 1
+                )[-1]
+            )
+        else:
+            func_name = "_list_outputs"
+            body = _strip_doc_string(
+                inspect.getsource(self.nipype_interface._list_outputs).split("\n", 1)[
+                    -1
+                ]
+            )
         body = self._process_inputs(body)
         body = self._misc_cleanups(body)
 
         if not body:
             return ""
         body = self.unwrap_nested_methods(
-            body,
-            additional_args=CALLABLES_ARGS,
+            body, additional_args=CALLABLES_ARGS, inputs_as_dict=True
         )
         body = self.replace_supers(body)
 
-        code_str = f"""def _list_outputs(inputs=None, stdout=None, stderr=None, output_dir=None):
-    parsed_inputs = _parse_inputs(inputs) if inputs else {{}}
+        code_str = f"""def {func_name}(inputs=None, stdout=None, stderr=None, output_dir=None):{self.parse_inputs_call}
 {body}
 """
         # Create separate function for each output field in the "callables" section
         for output_field in self.callable_output_fields:
             output_name = output_field[0]
             code_str += (
                 f"\n\n\ndef {output_name}_callable(output_dir, inputs, stdout, stderr):\n"
-                "    outputs = _list_outputs(output_dir=output_dir, inputs=inputs, stdout=stdout, stderr=stderr)\n"
+                f"    outputs = {func_name}(output_dir=output_dir, inputs=inputs, stdout=stdout, stderr=stderr)\n"
                 '    return outputs["' + output_name + '"]\n\n'
             )
         return code_str
@@ -419,9 +425,15 @@ def _process_inputs(self, body: str) -> str:
                 self.task_name,
             )
         body = input_re.sub(r"inputs['\1']", body)
-        body = re.sub(r"self\.(?!inputs)(\w+)", r"parsed_inputs['\1']", body)
+        body = re.sub(r"self\.(?!inputs)(\w+)\b(?!\()", r"parsed_inputs['\1']", body)
         return body
 
+    @property
+    def parse_inputs_call(self):
+        if not self.parse_inputs_code:
+            return ""
+        return "\n    _parse_inputs(inputs) if inputs else {}"
+
 
 def _strip_doc_string(body: str) -> str:
     if re.match(r"\s*(\"|')", body):