dropped StateIndex object in favour of simple int index

Author: tclose

pydra/engine/core.py

@@ -89,7 +89,7 @@ class Task(ty.Generic[DefType]):
     definition: DefType
     submitter: "Submitter | None"
     environment: "Environment | None"
-    state_index: state.StateIndex
+    state_index: int
     bindings: dict[str, ty.Any] | None = None  # Bindings for the task environment
 
     _inputs: dict[str, ty.Any] | None = None
@@ -100,7 +100,7 @@ def __init__(
         submitter: "Submitter",
         name: str,
         environment: "Environment | None" = None,
-        state_index: "state.StateIndex | None" = None,
+        state_index: int | None = None,
         hooks: TaskHooks | None = None,
     ):
         """
@@ -120,9 +120,6 @@ def __init__(
             4. Two or more concurrent new processes get to start
         """
 
-        if state_index is None:
-            state_index = state.StateIndex()
-
         if not isinstance(definition, TaskDef):
             raise ValueError(
                 f"Task definition ({definition!r}) must be a TaskDef, not {type(definition)}"

pydra/engine/lazy.py

@@ -9,7 +9,6 @@
     from .submitter import DiGraph, NodeExecution
     from .core import Task, Workflow
     from .specs import TaskDef
-    from .state import StateIndex
 
 
 T = ty.TypeVar("T")
@@ -49,7 +48,7 @@ def _get_value(
         self,
         workflow: "Workflow",
         graph: "DiGraph[NodeExecution]",
-        state_index: "StateIndex | None" = None,
+        state_index: int | None = None,
     ) -> ty.Any:
         """Return the value of a lazy field.
 
@@ -59,7 +58,7 @@ def _get_value(
             the workflow object
         graph: DiGraph[NodeExecution]
             the graph representing the execution state of the workflow
-        state_index : StateIndex, optional
+        state_index : int, optional
             the state index of the field to access
 
         Returns
@@ -95,7 +94,7 @@ def _get_value(
         self,
         workflow: "Workflow",
         graph: "DiGraph[NodeExecution]",
-        state_index: "StateIndex | None" = None,
+        state_index: int | None = None,
     ) -> ty.Any:
         """Return the value of a lazy field.
 
@@ -105,7 +104,7 @@ def _get_value(
             the workflow object
         graph: DiGraph[NodeExecution]
             the graph representing the execution state of the workflow
-        state_index : StateIndex, optional
+        state_index : int, optional
             the state index of the field to access
 
         Returns
@@ -134,7 +133,7 @@ def _get_value(
         self,
         workflow: "Workflow",
         graph: "DiGraph[NodeExecution]",
-        state_index: "StateIndex | None" = None,
+        state_index: int | None = None,
     ) -> ty.Any:
         """Return the value of a lazy field.
 
@@ -144,18 +143,14 @@ def _get_value(
             the workflow object
         graph: DiGraph[NodeExecution]
             the graph representing the execution state of the workflow
-        state_index : StateIndex, optional
+        state_index : int, optional
             the state index of the field to access
 
         Returns
         -------
         value : Any
             the resolved value of the lazy-field
         """
-        from pydra.engine.state import StateIndex
-
-        if state_index is None:
-            state_index = StateIndex()
 
         jobs = graph.node(self._node.name).matching_jobs(state_index)
 

pydra/engine/state.py

@@ -2,7 +2,6 @@
 
 from copy import deepcopy
 import itertools
-from collections import OrderedDict
 from functools import reduce
 import typing as ty
 from . import helpers_state as hlpst
@@ -15,117 +14,6 @@
 OutputsType = ty.TypeVar("OutputsType")
 
 
-class StateIndex:
-    """The collection of state indices that identifies a single element within the list
-    of tasks generated from a node
-
-    Parameters
-    ----------
-    indices : dict[str, int]
-        a dictionary of indices for each input field
-    """
-
-    indices: OrderedDict[str, int]
-
-    def __init__(
-        self, indices: dict[str, int] | ty.Sequence[tuple[str, int]] | None = None
-    ):
-        # We used ordered dict here to ensure the keys are always in the same order
-        # while OrderedDict is not strictly necessary for CPython 3.7+, we use it to
-        # signal that the order of the keys is important
-        if indices is None:
-            self.indices = OrderedDict()
-        else:
-            if isinstance(indices, dict):
-                indices = indices.items()
-            self.indices = OrderedDict(sorted(indices))
-
-    def __len__(self) -> int:
-        return len(self.indices)
-
-    def __iter__(self) -> ty.Generator[str, None, None]:
-        return iter(self.indices)
-
-    def __getitem__(self, key: str) -> int:
-        return self.indices[key]
-
-    def __lt__(self, other: "StateIndex") -> bool:
-        if list(self.indices) != list(other.indices):
-            raise ValueError(
-                f"StateIndex {self} does not contain the same indices in the same order "
-                f"as {other}: {list(self.indices)} != {list(other.indices)}"
-            )
-        return tuple(self.indices.items()) < tuple(other.indices.items())
-
-    def __repr__(self) -> str:
-        return (
-            "StateIndex(" + ", ".join(f"{n}={v}" for n, v in self.indices.items()) + ")"
-        )
-
-    def __hash__(self):
-        return hash(tuple(self.indices.items()))
-
-    def __eq__(self, other) -> bool:
-        return self.indices == other.indices
-
-    def __str__(self) -> str:
-        return "__".join(f"{n}-{i}" for n, i in self.indices.items())
-
-    def __bool__(self) -> bool:
-        return bool(self.indices)
-
-    def subset(self, state_names: ty.Iterable[str]) -> ty.Self:
-        """Create a new StateIndex with only the specified fields
-
-        Parameters
-        ----------
-        fields : list[str]
-            the fields to keep in the new StateIndex
-
-        Returns
-        -------
-        StateIndex
-            a new StateIndex with only the specified fields
-        """
-        return type(self)({k: v for k, v in self.indices.items() if k in state_names})
-
-    def missing(self, state_names: ty.Iterable[str]) -> ty.List[str]:
-        """Return the fields that are missing from the StateIndex
-
-        Parameters
-        ----------
-        fields : list[str]
-            the fields to check for
-
-        Returns
-        -------
-        list[str]
-            the fields that are missing from the StateIndex
-        """
-        return [f for f in state_names if f not in self.indices]
-
-    def matches(self, other: "StateIndex") -> bool:
-        """Check if the indices that are present in the other StateIndex match
-
-        Parameters
-        ----------
-        other : StateIndex
-            the other StateIndex to compare against
-
-        Returns
-        -------
-        bool
-            True if all the indices in the other StateIndex match
-        """
-        if isinstance(other, dict):
-            other = StateIndex(other)
-        if not set(self.indices).issuperset(other.indices):
-            raise ValueError(
-                f"StateIndex {self} does not contain all the indices in {other}"
-            )
-        return all(self.indices[k] == v for k, v in other.indices.items())
-
-
 class State:
     """
     A class that specifies a State of all tasks.
@@ -1314,7 +1202,7 @@ def _single_op_splits(self, op_single):
             keys = [op_single]
             return val, keys
 
-    def _get_element(self, value: ty.Any, field_name: str, ind: int):
+    def _get_element(self, value: ty.Any, field_name: str, ind: int) -> ty.Any:
         """
         Extracting element of the inputs taking into account
         container dimension of the specific element that can be set in self.state.cont_dim.
@@ -1329,6 +1217,11 @@ def _get_element(self, value: ty.Any, field_name: str, ind: int):
             name of the input field
         ind : int
             index of the element
+
+        Returns
+        -------
+        Any
+            specific element of the input field
         """
         if f"{self.name}.{field_name}" in self.cont_dim:
             return list(