Merge pull request #315 from djarecka/fix/ext_template

[mnt, bug] refactoring the template formatting (closes #314)

Author: djarecka

pydra/engine/core.py

@@ -423,9 +423,7 @@ def _collect_outputs(self):
         self.output_spec = output_from_inputfields(self.output_spec, self.inputs)
         output_klass = make_klass(self.output_spec)
         output = output_klass(**{f.name: None for f in attr.fields(output_klass)})
-        other_output = output.collect_additional_outputs(
-            self.input_spec, self.inputs, self.output_dir
-        )
+        other_output = output.collect_additional_outputs(self.inputs, self.output_dir)
         return attr.evolve(output, **run_output, **other_output)
 
     def split(self, splitter, overwrite=False, **kwargs):

pydra/engine/helpers.py

@@ -10,7 +10,7 @@
 from hashlib import sha256
 import subprocess as sp
 import getpass
-import uuid
+import re
 from time import strftime
 from traceback import format_exception
 
@@ -636,3 +636,36 @@ def position_adjustment(pos_args):
         cmd_args += el[1]
 
     return cmd_args
+
+
+def argstr_formatting(argstr, inputs, value_updates=None):
+    """ formatting argstr that have form {field_name},
+    using values from inputs and updating with value_update if provided
+    """
+    inputs_dict = attr.asdict(inputs)
+    # if there is a value that has to be updated (e.g. single value from a list)
+    if value_updates:
+        inputs_dict.update(value_updates)
+    # getting all fields that should be formatted, i.e. {field_name}, ...
+    inp_fields = re.findall("{\w+}", argstr)
+    val_dict = {}
+    for fld in inp_fields:
+        fld_name = fld[1:-1]  # extracting the name form {field_name}
+        fld_value = inputs_dict[fld_name]
+        if fld_value is attr.NOTHING:
+            # if value is NOTHING, nothing should be added to the command
+            val_dict[fld_name] = ""
+        else:
+            val_dict[fld_name] = fld_value
+
+    # formatting string based on the val_dict
+    argstr_formatted = argstr.format(**val_dict)
+    # removing extra commas and spaces after removing the field that have NOTHING
+    argstr_formatted = (
+        argstr_formatted.replace("[ ", "[")
+        .replace(" ]", "]")
+        .replace("[,", "[")
+        .replace(",]", "]")
+        .strip()
+    )
+    return argstr_formatted

pydra/engine/helpers_file.py

@@ -521,43 +521,93 @@ def template_update(inputs, map_copyfiles=None):
                 f"fields with output_file_template"
                 "has to be a string or Union[str, bool]"
             )
-        inp_val_set = getattr(inputs, fld.name)
+        dict_[fld.name] = template_update_single(field=fld, inputs_dict=dict_)
+    # using is and  == so it covers list and numpy arrays
+    updated_templ_dict = {
+        k: v
+        for k, v in dict_.items()
+        if not (getattr(inputs, k) is v or getattr(inputs, k) == v)
+    }
+    return updated_templ_dict
+
+
+def template_update_single(field, inputs_dict, spec_type="input"):
+    """Update a single template from the input_spec or output_spec
+    based on the value from inputs_dict
+    (checking the types of the fields, that have "output_file_template)"
+    """
+    from .specs import File
+
+    if spec_type == "input":
+        if field.type not in [str, ty.Union[str, bool]]:
+            raise Exception(
+                f"fields with output_file_template"
+                "has to be a string or Union[str, bool]"
+            )
+        inp_val_set = inputs_dict[field.name]
         if inp_val_set is not attr.NOTHING and not isinstance(inp_val_set, (str, bool)):
-            raise Exception(f"{fld.name} has to be str or bool, but {inp_val_set} set")
-        if isinstance(inp_val_set, bool) and fld.type is str:
             raise Exception(
-                f"type of {fld.name} is str, consider using Union[str, bool]"
+                f"{field.name} has to be str or bool, but {inp_val_set} set"
             )
+        if isinstance(inp_val_set, bool) and field.type is str:
+            raise Exception(
+                f"type of {field.name} is str, consider using Union[str, bool]"
+            )
+    elif spec_type == "output":
+        if field.type is not File:
+            raise Exception(
+                f"output {field.name} should be a File, but {field.type} set as the type"
+            )
+    else:
+        raise Exception(f"spec_type can be input or output, but {spec_type} provided")
+    if spec_type == "input" and isinstance(inputs_dict[field.name], str):
+        return inputs_dict[field.name]
+    elif spec_type == "input" and inputs_dict[field.name] is False:
+        # if input fld is set to False, the fld shouldn't be used (setting NOTHING)
+        return attr.NOTHING
+    else:  # inputs_dict[field.name] is True or spec_type is output
+        template = field.metadata["output_file_template"]
+        # as default, we assume that keep_extension is True
+        keep_extension = field.metadata.get("keep_extension", True)
+        value = _template_formatting(
+            template, inputs_dict, keep_extension=keep_extension
+        )
+        return value
 
-        if isinstance(inp_val_set, str):
-            dict_[fld.name] = inp_val_set
-        elif inp_val_set is False:
-            # if False, the field should not be used, so setting attr.NOTHING
-            dict_[fld.name] = attr.NOTHING
-        else:  # True or attr.NOTHING
-            template = fld.metadata["output_file_template"]
-            value = template.format(**dict_)
-            value = removing_nothing(value)
-            dict_[fld.name] = value
-    return {k: v for k, v in dict_.items() if getattr(inputs, k) is not v}
-
-
-def removing_nothing(template_str):
-    """ removing all fields that had NOTHING"""
-    if "NOTHING" not in template_str:
-        return template_str
-    regex = re.compile(r"[^a-zA-Z_\-]")
-    fields_str = regex.sub(" ", template_str)
-    for fld in fields_str.split():
-        if "NOTHING" in fld:
-            template_str = template_str.replace(fld, "")
-    return (
-        template_str.replace("[ ", "[")
-        .replace(" ]", "]")
-        .replace(",]", "]")
-        .replace("[,", "[")
-        .strip()
-    )
+
+def _template_formatting(template, inputs_dict, keep_extension=True):
+    """Formatting a single template based on values from inputs_dict.
+    Taking into account that field values and template could have file extensions
+    (assuming that if template has extension, the field value extension is removed,
+    if field has extension, and no template extension, than it is moved to the end),
+    """
+    inp_fields = re.findall("{\w+}", template)
+    if len(inp_fields) == 0:
+        return template
+    elif len(inp_fields) == 1:
+        fld_name = inp_fields[0][1:-1]
+        fld_value = inputs_dict[fld_name]
+        if fld_value is attr.NOTHING:
+            return attr.NOTHING
+        fld_value = str(fld_value)  # in case it's a path
+        filename, *ext = fld_value.split(".", maxsplit=1)
+        # if keep_extension is False, the extensions are removed
+        if keep_extension is False:
+            ext = []
+        if template.endswith(inp_fields[0]):
+            # if no suffix added in template, the simplest formatting should work
+            # recreating fld_value with the updated extension
+            fld_value_upd = ".".join([filename] + ext)
+            formatted_value = template.format(**{fld_name: fld_value_upd})
+        elif "." not in template:  # the template doesn't have its own extension
+            # if the fld_value has extension, it will be moved to the end
+            formatted_value = ".".join([template.format(**{fld_name: filename})] + ext)
+        else:  # template has its own extension
+            # removing fld_value extension if any
+            formatted_value = template.format(**{fld_name: filename})
+        return formatted_value
+    else:
+        raise NotImplementedError("should we allow for more args in the template?")
 
 
 def is_local_file(f):

pydra/engine/specs.py

@@ -3,6 +3,8 @@
 from pathlib import Path
 import typing as ty
 
+from .helpers_file import template_update_single
+
 
 def attr_fields(x):
     return x.__attrs_attrs__
@@ -33,7 +35,7 @@ class SpecInfo:
 class BaseSpec:
     """The base dataclass specs for all inputs and outputs."""
 
-    def collect_additional_outputs(self, input_spec, inputs, output_dir):
+    def collect_additional_outputs(self, inputs, output_dir):
         """Get additional outputs."""
         return {}
 
@@ -213,6 +215,7 @@ def check_metadata(self):
             "position",
             "requires",
             "separate_ext",
+            "keep_extension",
             "xor",
             "sep",
         }
@@ -322,7 +325,7 @@ class ShellOutSpec(BaseSpec):
     stderr: ty.Union[File, str]
     """The process' standard input."""
 
-    def collect_additional_outputs(self, input_spec, inputs, output_dir):
+    def collect_additional_outputs(self, inputs, output_dir):
         """Collect additional outputs from shelltask output_spec."""
         additional_out = {}
         for fld in attr_fields(self):
@@ -379,24 +382,13 @@ def _field_metadata(self, fld, inputs, output_dir):
         """Collect output file if metadata specified."""
         if "value" in fld.metadata:
             return output_dir / fld.metadata["value"]
+        # this block is only run if "output_file_template" is provided in output_spec
+        # if the field is set in input_spec with output_file_template,
+        # than the field already should have value
         elif "output_file_template" in fld.metadata:
-            sfx_tmpl = (output_dir / fld.metadata["output_file_template"]).suffixes
-            if sfx_tmpl:
-                # removing suffix from input field if template has it's own suffix
-                inputs_templ = {
-                    k: v.split(".")[0]
-                    for k, v in inputs.__dict__.items()
-                    if isinstance(v, str)
-                }
-            else:
-                inputs_templ = {
-                    k: v for k, v in inputs.__dict__.items() if isinstance(v, str)
-                }
-            out_path = output_dir / fld.metadata["output_file_template"].format(
-                **inputs_templ
-            )
-            return out_path
-
+            inputs_templ = attr.asdict(inputs)
+            value = template_update_single(fld, inputs_templ, spec_type="output")
+            return output_dir / value
         elif "callable" in fld.metadata:
             return fld.metadata["callable"](fld.name, output_dir)
         else:

pydra/engine/task.py

@@ -57,8 +57,8 @@
     SingularitySpec,
     attr_fields,
 )
-from .helpers import ensure_list, execute, position_adjustment
-from .helpers_file import template_update, is_local_file, removing_nothing
+from .helpers import ensure_list, execute, position_adjustment, argstr_formatting
+from .helpers_file import template_update, is_local_file
 
 
 class FunctionTask(TaskBase):
@@ -431,38 +431,40 @@ def _command_pos_args(self, field, state_ind, ind):
 
         cmd_add = []
         if field.type is bool:
+            # if value is simply True the original argstr is used,
+            # if False, nothing is added to the command
             if value is True:
                 cmd_add.append(argstr)
         else:
             sep = field.metadata.get("sep", " ")
             if argstr.endswith("...") and isinstance(value, list):
                 argstr = argstr.replace("...", "")
+                # if argstr has a more complex form, with "{input_field}"
                 if "{" in argstr and "}" in argstr:
                     argstr_formatted_l = []
                     for val in value:
-                        argstr_f = argstr.format(**{field.name: val}).format(
-                            **attr.asdict(self.inputs)
+                        argstr_f = argstr_formatting(
+                            argstr, self.inputs, value_updates={field.name: val}
                         )
-                        argstr_formatted_l.append(removing_nothing(argstr_f))
-
+                        argstr_formatted_l.append(argstr_f)
                     cmd_el_str = sep.join(argstr_formatted_l)
-                else:
+                else:  # argstr has a simple form, e.g. "-f", or "--f"
                     cmd_el_str = sep.join([f" {argstr} {val}" for val in value])
             else:
                 # in case there are ... when input is not a list
                 argstr = argstr.replace("...", "")
                 if isinstance(value, list):
                     cmd_el_str = sep.join([str(val) for val in value])
                     value = cmd_el_str
-
+                # if argstr has a more complex form, with "{input_field}"
                 if "{" in argstr and "}" in argstr:
-                    argstr_f = argstr.format(**attr.asdict(self.inputs))
-                    cmd_el_str = removing_nothing(argstr_f)
-                else:
+                    cmd_el_str = argstr_formatting(argstr, self.inputs)
+                else:  # argstr has a simple form, e.g. "-f", or "--f"
                     if value:
                         cmd_el_str = f"{argstr} {value}"
                     else:
                         cmd_el_str = ""
+            # removing double spacing
             cmd_el_str = cmd_el_str.strip().replace("  ", " ")
             if cmd_el_str:
                 cmd_add += cmd_el_str.split(" ")