reimplementing lazy-field split

Author: tclose

pydra/engine/lazy.py

@@ -68,6 +68,12 @@ def _get_value(
         """
         raise NotImplementedError("LazyField is an abstract class")
 
+    def split(self) -> "LazyField":
+        """ "Splits" the lazy field over an array of nodes by replacing the sequence type
+        of the lazy field with StateArray to signify that it will be "split" across
+        """
+        raise NotImplementedError("LazyField is an abstract class")
+
 
 @attrs.define(kw_only=True)
 class LazyInField(LazyField[T]):
@@ -116,6 +122,25 @@ def _get_value(
         value = self._apply_cast(value)
         return value
 
+    def split(self) -> "LazyField":
+        """ "Splits" the lazy field over an array of nodes by replacing the sequence type
+        of the lazy field with StateArray to signify that it will be "split" across
+        """
+        from ..utils.typing import TypeParser  # pylint: disable=import-outside-toplevel
+
+        assert not isinstance(self, LazyInField)
+
+        if not TypeParser.matches_type(self.type, list):
+            raise TypeError(
+                f"Cannot split non-sequence field {self}  of type {self.type}"
+            )
+
+        return type(self)(
+            name=self.name,
+            field=self.field,
+            type=StateArray[TypeParser.get_item_type(self.type)],
+        )
+
 
 @attrs.define(kw_only=True)
 class LazyOutField(LazyField[T]):
@@ -206,6 +231,34 @@ def retrieve_from_job(job: "Task[DefType]") -> ty.Any:
         # else:
         return [retrieve_from_job(j) for j in jobs]
 
+    def split(self) -> "LazyField":
+        """ "Splits" the lazy field over an array of nodes by replacing the sequence type
+        of the lazy field with StateArray to signify that it will be "split" across
+        """
+        from ..utils.typing import TypeParser  # pylint: disable=import-outside-toplevel
+
+        # Modify the type of the lazy field to include the split across a state-array
+        inner_type, prev_split_depth = TypeParser.strip_splits(self.type)
+        assert prev_split_depth <= 1
+        if inner_type is ty.Any:
+            type_ = StateArray[ty.Any]
+        elif TypeParser.matches_type(inner_type, list):
+            item_type = TypeParser.get_item_type(inner_type)
+            type_ = StateArray[item_type]
+        else:
+            raise TypeError(
+                f"Cannot split non-sequence field {self}  of type {inner_type}"
+            )
+        if prev_split_depth:
+            type_ = StateArray[
+                type_
+            ]  # FIXME: This nesting of StateArray is probably unnecessary
+        return type(self)[type_](
+            name=self.name,
+            field=self.field,
+            type=type_,
+        )
+
     @property
     def _source(self):
         return self._node