fixed handling of super methods, added supported for aggregate_outputs and fix bug with inputs handling in _list_outputs

Author: tclose

nipype2pydra/interface/base.py

@@ -31,6 +31,8 @@
     cleanup_function_body,
     insert_args_in_signature,
     extract_args,
+    strip_comments,
+    find_super_method,
 )
 from ..statements import (
     ImportStatement,
@@ -1021,6 +1023,11 @@ def _misc_cleanups(self, body: str) -> str:
         )
         body = re.sub(r"\bruntime\.(stdout|stderr)", r"\1", body)
         body = re.sub(r"\boutputs\.(\w+)", r"outputs['\1']", body)
+        body = re.sub(r"getattr\(inputs, ([^)]+)\)", r"inputs[\1]", body)
+        body = re.sub(
+            r"setattr\(outputs, ([^,]+), ([^)]+)\)", r"outputs[\1] = \2", body
+        )
+        body = body.replace("TraitError", "KeyError")
         body = body.replace("os.getcwd()", "output_dir")
         return body
 
@@ -1209,7 +1216,9 @@ def process_method(
         if method.__name__ in self.method_args:
             args += [
                 f"{a}=None"
-                for a in (list(self.method_args[method.__name__]) + additional_args)
+                for a in (
+                    list(self.method_args[method.__name__]) + list(additional_args)
+                )
             ]
         # Insert method args in signature if present
         return_types, method_body = post.split(":", maxsplit=1)
@@ -1291,7 +1300,9 @@ def replace_supers(self, method_body, super_base=None):
         def replace_super(match):
             super_method, base = find_super_method(super_base, match.group(1))
             try:
-                return self.SPECIAL_SUPER_MAPPINGS[super_method]
+                return self.SPECIAL_SUPER_MAPPINGS[super_method].format(
+                    args=match.group(2)
+                )
             except KeyError:
                 try:
                     return name_map[match.group(1)] + "(" + match.group(2) + ")"
@@ -1316,7 +1327,9 @@ def unwrap_nested_methods(
         """
         # 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]
+        method_names = [m.__name__ for m in self.referenced_methods] + list(
+            self.INCLUDED_METHODS
+        )
         method_body = strip_comments(method_body)
         omitted_methods = {}
         for method_name in set(
@@ -1329,8 +1342,11 @@ def unwrap_nested_methods(
         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:]
+                args, post = extract_args(args)[1:]
+                new_body += self.SPECIAL_SUPER_MAPPINGS[omitted_methods[name]].format(
+                    args=", ".join(args)
+                )
+                new_body += post[1:]  # drop the leading parenthesis
                 continue
             # Assign additional return values (which were previously saved to member
             # attributes) to new variables from the method call
@@ -1375,8 +1391,9 @@ def unwrap_nested_methods(
         return cleanup_function_body(method_body)
 
     SPECIAL_SUPER_MAPPINGS = {
-        CommandLine._list_outputs: "{}",
+        CommandLine._list_outputs: "{{}}",
         CommandLine._format_arg: "argstr.format(**inputs)",
+        CommandLine._filename_from_source: "{args} + '_generated'",
         BaseInterface._check_version_requirements: "[]",
     }
 
@@ -1431,19 +1448,3 @@ def pytest_configure(config):
 else:
     CATCH_CLI_EXCEPTIONS = True
 """
-
-
-def find_super_method(
-    super_base: type, method_name: str
-) -> ty.Tuple[ty.Callable, type]:
-    for base in super_base.__mro__[1:]:
-        if method_name in base.__dict__:  # Found the match
-            return getattr(base, method_name), base
-    raise RuntimeError(
-        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

@@ -13,6 +13,7 @@
     UsedSymbols,
     split_source_into_statements,
     INBUILT_NIPYPE_TRAIT_NAMES,
+    find_super_method,
 )
 from fileformats.core.mixin import WithClassifiers
 from fileformats.generic import File, Directory
@@ -243,6 +244,10 @@ def callable_output_fields(self):
             )
         ]
 
+    @property
+    def callable_output_field_names(self):
+        return [f[0] for f in self.callable_output_fields]
+
     @cached_property
     def _format_arg_body(self):
         if "_format_arg" not in self.nipype_interface.__dict__:
@@ -295,7 +300,7 @@ def format_arg_code(self):
     if {val_arg} is None:
         return ""
 {body}
-
+    return argstr.format(**inputs)
 
 """
         for field_name in self.formatted_input_field_names:
@@ -326,7 +331,7 @@ def parse_inputs_code(self) -> str:
         body = self.unwrap_nested_methods(body, inputs_as_dict=True)
         body = self.replace_supers(body)
 
-        code_str = "def _parse_inputs(inputs):\n    parsed_inputs = {{}}"
+        code_str = "def _parse_inputs(inputs):\n    parsed_inputs = {}"
         if re.findall(r"\bargstrs\b", body):
             code_str += f"\n    argstrs = {self._format_argstrs!r}"
         code_str += f"""
@@ -374,40 +379,70 @@ def callables_code(self):
 
         if not self.callable_output_fields:
             return ""
-        if hasattr(self.nipype_interface, "aggregate_outputs"):
+        code_str = ""
+        if (
+            find_super_method(self.nipype_interface, "aggregate_outputs")[1]
+            is not BaseInterface
+        ):
             func_name = "aggregate_outputs"
             body = _strip_doc_string(
                 inspect.getsource(self.nipype_interface.aggregate_outputs).split(
                     "\n", 1
                 )[-1]
             )
+            need_list_outputs = bool(re.findall(r"\b_list_outputs\b", body))
+            body = self._process_inputs(body)
+            body = self._misc_cleanups(body)
+
+            if not body:
+                return ""
+            body = self.unwrap_nested_methods(
+                body, additional_args=CALLABLES_ARGS, inputs_as_dict=True
+            )
+            body = self.replace_supers(body)
+
+            code_str += f"""def aggregate_outputs(inputs=None, stdout=None, stderr=None, output_dir=None):
+    inputs = attrs.asdict(inputs){self.parse_inputs_call}
+    needed_outputs = {self.callable_output_field_names!r}
+{body}
+
+
+"""
+            inputs_as_dict_call = ""
+
         else:
             func_name = "_list_outputs"
+            inputs_as_dict_call = "\n    inputs = attrs.asdict(inputs)"
+            need_list_outputs = True
+
+        if need_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)
+            body = self._process_inputs(body)
+            body = self._misc_cleanups(body)
 
-        if not body:
-            return ""
-        body = self.unwrap_nested_methods(
-            body, additional_args=CALLABLES_ARGS, inputs_as_dict=True
-        )
-        body = self.replace_supers(body)
+            if not body:
+                return ""
+            body = self.unwrap_nested_methods(
+                body, additional_args=CALLABLES_ARGS, inputs_as_dict=True
+            )
+            body = self.replace_supers(body)
 
-        code_str = f"""def {func_name}(inputs=None, stdout=None, stderr=None, output_dir=None):{self.parse_inputs_call}
+        code_str += f"""def _list_outputs(inputs=None, stdout=None, stderr=None, output_dir=None):{inputs_as_dict_call}{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"
                 f"    outputs = {func_name}(output_dir=output_dir, inputs=inputs, stdout=stdout, stderr=stderr)\n"
-                '    return outputs["' + output_name + '"]\n\n'
+                '    return outputs.get("' + output_name + '", attrs.NOTHING)\n\n'
             )
         return code_str
 
@@ -432,7 +467,7 @@ def _process_inputs(self, body: str) -> str:
     def parse_inputs_call(self):
         if not self.parse_inputs_code:
             return ""
-        return "\n    _parse_inputs(inputs) if inputs else {}"
+        return "\n    parsed_inputs = _parse_inputs(inputs) if inputs else {}"
 
 
 def _strip_doc_string(body: str) -> str:

nipype2pydra/utils/__init__.py

@@ -20,6 +20,8 @@
     types_converter,
     unwrap_nested_type,
     get_return_line,
+    find_super_method,
+    strip_comments,
     INBUILT_NIPYPE_TRAIT_NAMES,
 )
 from .symbols import (  # noqa: F401

nipype2pydra/utils/misc.py

@@ -539,3 +539,19 @@ def get_return_line(func: ty.Union[str, ty.Callable]) -> str:
     if not match:
         return None
     return match.group(1).strip()
+
+
+def find_super_method(
+    super_base: type, method_name: str
+) -> ty.Tuple[ty.Callable, type]:
+    for base in super_base.__mro__[1:]:
+        if method_name in base.__dict__:  # Found the match
+            return getattr(base, method_name), base
+    raise RuntimeError(
+        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)