moved xor into *.define decorators from *.arg fields

tclose · tclose · commit 9df11910b81c · 2025-03-07T17:45:22.000+11:00
diff --git a/pydra/design/base.py b/pydra/design/base.py
@@ -248,10 +248,6 @@ class Arg(Field):
         Names of the inputs that are required together with the field.
     allowed_values: Sequence, optional
         List of allowed values for the field.
-    xor: Sequence[str | None], optional
-        Names of args that are exclusive mutually exclusive, which must include
-        the name of the current field. If this list includes None, then none of the
-        fields need to be set.
     copy_mode: File.CopyMode, optional
         The mode of copying the file, by default it is File.CopyMode.any
     copy_collation: File.CopyCollation, optional
@@ -266,25 +262,11 @@ class Arg(Field):
     """
 
     allowed_values: frozenset = attrs.field(default=(), converter=frozenset)
-    xor: frozenset[str | None] = attrs.field(default=(), converter=frozenset)
     copy_mode: File.CopyMode = File.CopyMode.any
     copy_collation: File.CopyCollation = File.CopyCollation.any
     copy_ext_decomp: File.ExtensionDecomposition = File.ExtensionDecomposition.single
     readonly: bool = False
 
-    @xor.validator
-    def _xor_validator(self, _, value):
-        for v in value:
-            if not isinstance(v, (str, type(None))):
-                raise ValueError(
-                    f"xor values must be strings or None, not {v} ({self!r})"
-                )
-        if value and self.type not in (ty.Any, bool) and not is_optional(self.type):
-            raise ValueError(
-                f"Fields that have 'xor' must be of boolean or optional type, "
-                f"not type {self.type} ({self!r})"
-            )
-
 
 @attrs.define(kw_only=True, slots=False)
 class Out(Field):
@@ -418,6 +400,7 @@ def make_task_def(
     name: str | None = None,
     bases: ty.Sequence[type] = (),
     outputs_bases: ty.Sequence[type] = (),
+    xor: ty.Sequence[str | None] | ty.Sequence[ty.Sequence[str | None]] = (),
 ):
     """Create a task definition class and its outputs definition class from the
     input and output fields provided to the decorator/function.
@@ -442,14 +425,26 @@ def make_task_def(
         The base classes for the task definition class, by default ()
     outputs_bases : ty.Sequence[type], optional
         The base classes for the outputs definition class, by default ()
+    xor: Sequence[str | None] | Sequence[Sequence[str | None]], optional
+        Names of args that are exclusive mutually exclusive, which must include
+        the name of the current field. If this list includes None, then none of the
+        fields need to be set.
 
     Returns
     -------
     klass : type
         The class created using the attrs package
     """
 
-    spec_type._check_arg_refs(inputs, outputs)
+    # Convert a single xor set into a set of xor sets
+    if not xor:
+        xor = frozenset()
+    elif all(isinstance(x, str) or x is None for x in xor):
+        xor = frozenset([frozenset(xor)])
+    else:
+        xor = frozenset(frozenset(x) for x in xor)
+
+    spec_type._check_arg_refs(inputs, outputs, xor)
 
     # Check that the field attributes are valid after all fields have been set
     # (especially the type)
@@ -521,6 +516,8 @@ def make_task_def(
                 **attrs_kwargs,
             ),
         )
+        # Store the xor sets for the class
+        klass._xor = xor
         klass.__annotations__[arg.name] = field_type
 
     # Create class using attrs package, will create attributes for all columns and
diff --git a/pydra/design/boutiques.py b/pydra/design/boutiques.py
@@ -28,10 +28,6 @@ class arg(shell.arg):
         List of allowed values for the field.
     requires: list, optional
         Names of the inputs that are required together with the field.
-    xor: list[str | None], optional
-        Names of args that are exclusive mutually exclusive, which must include
-        the name of the current field. If this list includes None, then none of the
-        fields need to be set.
     copy_mode: File.CopyMode, optional
         The mode of copying the file, by default it is File.CopyMode.any
     copy_collation: File.CopyCollation, optional
diff --git a/pydra/design/python.py b/pydra/design/python.py
@@ -33,10 +33,6 @@ class arg(Arg):
         List of allowed values for the field.
     requires: list, optional
         Names of the inputs that are required together with the field.
-    xor: list[str | None], optional
-        Names of args that are exclusive mutually exclusive, which must include
-        the name of the current field. If this list includes None, then none of the
-        fields need to be set.
     copy_mode: File.CopyMode, optional
         The mode of copying the file, by default it is File.CopyMode.any
     copy_collation: File.CopyCollation, optional
@@ -106,6 +102,7 @@ def define(
     bases: ty.Sequence[type] = (),
     outputs_bases: ty.Sequence[type] = (),
     auto_attribs: bool = True,
+    xor: ty.Sequence[str | None] | ty.Sequence[ty.Sequence[str | None]] = (),
 ) -> "PythonDef":
     """
     Create an interface for a function or a class.
@@ -120,6 +117,10 @@ def define(
         The outputs of the function or class.
     auto_attribs : bool
         Whether to use auto_attribs mode when creating the class.
+    xor: Sequence[str | None] | Sequence[Sequence[str | None]], optional
+        Names of args that are exclusive mutually exclusive, which must include
+        the name of the current field. If this list includes None, then none of the
+        fields need to be set.
 
     Returns
     -------
@@ -172,6 +173,7 @@ def make(wrapped: ty.Callable | type) -> PythonDef:
             klass=klass,
             bases=bases,
             outputs_bases=outputs_bases,
+            xor=xor,
         )
 
         return defn
diff --git a/pydra/design/shell.py b/pydra/design/shell.py
@@ -56,10 +56,6 @@ class arg(Arg):
         List of allowed values for the field.
     requires: list, optional
         List of field names that are required together with the field.
-    xor: list[str | None], optional
-        Names of args that are exclusive mutually exclusive, which must include
-        the name of the current field. If this list includes None, then none of the
-        fields need to be set.
     copy_mode: File.CopyMode, optional
         The mode of copying the file, by default it is File.CopyMode.any
     copy_collation: File.CopyCollation, optional
@@ -194,10 +190,6 @@ class outarg(arg, Out):
         List of allowed values for the field.
     requires: list, optional
         List of field names that are required together with the field.
-    xor: list[str | None], optional
-        Names of args that are exclusive mutually exclusive, which must include
-        the name of the current field. If this list includes None, then none of the
-        fields need to be set.
     copy_mode: File.CopyMode, optional
         The mode of copying the file, by default it is File.CopyMode.any
     copy_collation: File.CopyCollation, optional
@@ -291,6 +283,7 @@ def define(
     outputs_bases: ty.Sequence[type] = (),
     auto_attribs: bool = True,
     name: str | None = None,
+    xor: ty.Sequence[str | None] | ty.Sequence[ty.Sequence[str | None]] = (),
 ) -> "ShellDef":
     """Create a task definition for a shell command. Can be used either as a decorator on
     the "canonical" dataclass-form of a task definition or as a function that takes a
@@ -337,6 +330,10 @@ def define(
         as they appear in the template
     name: str | None
         The name of the returned class
+    xor: Sequence[str | None] | Sequence[Sequence[str | None]], optional
+        Names of args that are exclusive mutually exclusive, which must include
+        the name of the current field. If this list includes None, then none of the
+        fields need to be set.
 
     Returns
     -------
@@ -446,6 +443,7 @@ def make(
             klass=klass,
             bases=bases,
             outputs_bases=outputs_bases,
+            xor=xor,
         )
         return defn
 
diff --git a/pydra/design/tests/test_shell.py b/pydra/design/tests/test_shell.py
@@ -469,7 +469,7 @@ def test_interface_template_with_type_overrides():
 def Ls(request):
     if request.param == "static":
 
-        @shell.define
+        @shell.define(xor=["complete_date", "date_format_str", None])
         class Ls(ShellDef["Ls.Outputs"]):
             executable = "ls"
 
@@ -502,14 +502,12 @@ class Ls(ShellDef["Ls.Outputs"]):
                 argstr="-T",
                 default=False,
                 requires=["long_format"],
-                xor=["complete_date", "date_format_str", None],
             )
             date_format_str: str | None = shell.arg(
                 help="format string for ",
                 argstr="-D",
                 default=None,
                 requires=["long_format"],
-                xor=["complete_date", "date_format_str", None],
             )
 
             @shell.outputs
diff --git a/pydra/design/workflow.py b/pydra/design/workflow.py
@@ -37,10 +37,6 @@ class arg(Arg):
         List of allowed values for the field.
     requires: list, optional
         Names of the inputs that are required together with the field.
-    xor: list[str | None], optional
-        Names of args that are exclusive mutually exclusive, which must include
-        the name of the current field. If this list includes None, then none of the
-        fields need to be set.
     copy_mode: File.CopyMode, optional
         The mode of copying the file, by default it is File.CopyMode.any
     copy_collation: File.CopyCollation, optional
@@ -111,6 +107,7 @@ def define(
     outputs_bases: ty.Sequence[type] = (),
     lazy: list[str] | None = None,
     auto_attribs: bool = True,
+    xor: ty.Sequence[str | None] | ty.Sequence[ty.Sequence[str | None]] = (),
 ) -> "WorkflowDef":
     """
     Create an interface for a function or a class. Can be used either as a decorator on
@@ -126,6 +123,10 @@ def define(
         The outputs of the function or class.
     auto_attribs : bool
         Whether to use auto_attribs mode when creating the class.
+    xor: Sequence[str | None] | Sequence[Sequence[str | None]], optional
+        Names of args that are exclusive mutually exclusive, which must include
+        the name of the current field. If this list includes None, then none of the
+        fields need to be set.
 
     Returns
     -------
@@ -183,6 +184,7 @@ def make(wrapped: ty.Callable | type) -> TaskDef:
             klass=klass,
             bases=bases,
             outputs_bases=outputs_bases,
+            xor=xor,
         )
 
         return defn
diff --git a/pydra/engine/specs.py b/pydra/engine/specs.py
@@ -173,6 +173,11 @@ def reset(self):
 class TaskDef(ty.Generic[OutputsType]):
     """Base class for all task definitions"""
 
+    # Class attributes
+    _xor: frozenset[frozenset[str | None]] = (
+        frozenset()
+    )  # overwritten in derived classes
+
     # The following fields are used to store split/combine state information
     _splitter = attrs.field(default=None, init=False, repr=False)
     _combiner = attrs.field(default=None, init=False, repr=False)
@@ -501,20 +506,6 @@ def _rule_violations(self) -> list[str]:
             ):
                 errors.append(f"Mandatory field {field.name!r} is not set")
 
-            # Collect alternative fields associated with this field.
-            if field.xor:
-                mutually_exclusive = {name: self[name] for name in field.xor if name}
-                are_set = [f"{n}={v!r}" for n, v in mutually_exclusive.items() if v]
-                if len(are_set) > 1:
-                    errors.append(
-                        f"Mutually exclusive fields ({', '.join(are_set)}) are set together"
-                    )
-                elif not are_set and None not in field.xor:
-                    errors.append(
-                        "At least one of the mutually exclusive fields should be set: "
-                        + ", ".join(f"{n}={v!r}" for n, v in mutually_exclusive.items())
-                    )
-
             # Raise error if any required field is unset.
             if (
                 not (
@@ -538,6 +529,19 @@ def _rule_violations(self) -> list[str]:
                 errors.append(
                     f"{field.name!r} requires{qualification} {[str(r) for r in field.requires]}"
                 )
+        # Collect alternative fields associated with this field.
+        for xor_set in self._xor:
+            mutually_exclusive = {name: self[name] for name in xor_set if name}
+            are_set = [f"{n}={v!r}" for n, v in mutually_exclusive.items() if v]
+            if len(are_set) > 1:
+                errors.append(
+                    f"Mutually exclusive fields ({', '.join(are_set)}) are set together"
+                )
+            elif not are_set and None not in xor_set:
+                errors.append(
+                    "At least one of the mutually exclusive fields should be set: "
+                    + ", ".join(f"{n}={v!r}" for n, v in mutually_exclusive.items())
+                )
         return errors
 
     def _check_rules(self):
@@ -552,7 +556,12 @@ def _check_rules(self):
             )
 
     @classmethod
-    def _check_arg_refs(cls, inputs: list[Arg], outputs: list[Out]) -> None:
+    def _check_arg_refs(
+        cls,
+        inputs: list[Arg],
+        outputs: list[Out],
+        xor: frozenset[frozenset[str | None]],
+    ) -> None:
         """
         Checks if all fields referenced in requirements and xor are present in the inputs
         are valid field names
@@ -567,12 +576,22 @@ def _check_arg_refs(cls, inputs: list[Arg], outputs: list[Out]) -> None:
                     "'Unrecognised' field names in referenced in the requirements "
                     f"of {field} " + str(list(unrecognised))
                 )
-        for inpt in inputs.values():
-            if unrecognised := inpt.xor - (input_names | {None}):
+
+        for xor_set in xor:
+            if unrecognised := xor_set - (input_names | {None}):
                 raise ValueError(
-                    "'Unrecognised' field names in referenced in the xor "
-                    f"of {inpt} " + str(list(unrecognised))
+                    f"'Unrecognised' field names in referenced in the xor {xor_set} "
+                    + str(list(unrecognised))
                 )
+            for field_name in xor_set:
+                if field_name is None:  # i.e. none of the fields being set is valid
+                    continue
+                type_ = inputs[field_name].type
+                if type_ not in (ty.Any, bool) and not is_optional(type_):
+                    raise ValueError(
+                        f"Fields included in a 'xor' ({field.name!r}) must be of boolean "
+                        f"or optional types, not type {type_}"
+                    )
 
     def _check_resolved(self):
         """Checks that all the fields in the definition have been resolved"""
@@ -762,6 +781,8 @@ def _from_task(cls, task: "Task[WorkflowDef]") -> Self:
         for name, lazy_field in attrs_values(workflow.outputs).items():
             try:
                 val_out = lazy_field._get_value(workflow=workflow, graph=exec_graph)
+                if isinstance(val_out, StateArray):
+                    val_out = list(val_out)  # implicitly combine state arrays
                 output_wf[name] = val_out
             except (ValueError, AttributeError):
                 output_wf[name] = None
diff --git a/pydra/engine/tests/test_shelltask.py b/pydra/engine/tests/test_shelltask.py
diff --git a/pydra/engine/tests/test_shelltask_inputspec.py b/pydra/engine/tests/test_shelltask_inputspec.py
