moving some functions from helper_state to the State methods; some mnt. edits

Author: djarecka

pydra/engine/helpers_state.py

@@ -776,110 +776,6 @@ def map_splits(split_iter, inputs, cont_dim=None):
         }
 
 
-# Functions for merging and completing splitters in states.
-def connect_splitters(splitter, other_states, state_fields=False):
-    """Link splitters."""
-    if splitter:
-        # if splitter is string, have to check if this is Left or Right part (Left is required)
-        if isinstance(splitter, str):
-            # so this is the Left part
-            if splitter.startswith("_"):
-                left_part = _complete_left(
-                    left=splitter, other_states=other_states, state_fields=state_fields
-                )
-                right_part = None
-            else:  # this is Right part
-                left_part = _complete_left(
-                    other_states=other_states, state_fields=state_fields
-                )
-                right_part = splitter
-        # if splitter is tuple, it has to be either Left or Right part
-        elif isinstance(splitter, tuple):
-            lr_flag = _left_right_check(splitter, other_states=other_states)
-            if lr_flag == "Left":
-                left_part = _complete_left(left=splitter, other_states=other_states)
-                right_part = None
-            elif lr_flag == "Right":
-                left_part = _complete_left(other_states=other_states)
-                right_part = splitter
-            else:
-                raise Exception("splitter mix Left and Right parts in scalar splitter")
-        elif isinstance(splitter, list):
-            lr_flag = _left_right_check(splitter, other_states=other_states)
-            if lr_flag == "Left":
-                left_part = _complete_left(
-                    left=splitter, other_states=other_states, state_fields=state_fields
-                )
-                right_part = None
-            elif lr_flag == "Right":
-                left_part = _complete_left(
-                    other_states=other_states, state_fields=state_fields
-                )
-                right_part = splitter
-            elif (
-                _left_right_check(splitter[0], other_states=other_states) == "Left"
-                and _left_right_check(splitter[1], other_states=other_states) == "Right"
-            ):
-                left_part = _complete_left(
-                    left=splitter[0],
-                    other_states=other_states,
-                    state_fields=state_fields,
-                )
-                right_part = splitter[1]
-            else:
-                raise Exception("splitter doesn't have separated Left and Right parts")
-        else:
-            raise Exception(
-                "splitter has to be str, tuple or list, "
-                "{} was provided".format(type(splitter))
-            )
-    else:
-        # if there is no splitter, I create the Left part
-        left_part = _complete_left(other_states=other_states, state_fields=state_fields)
-        right_part = None
-    if right_part:
-        splitter = [deepcopy(left_part), deepcopy(right_part)]
-    else:
-        splitter = deepcopy(left_part)
-    return splitter, left_part, right_part
-
-
-def _complete_left(other_states, left=None, state_fields=False):
-    """Add all splitters from previous nodes (completing left part)."""
-    if left:
-        rpn_left = splitter2rpn(
-            left, other_states=other_states, state_fields=state_fields
-        )
-        for name, (st, inp) in list(other_states.items())[::-1]:
-            if "_{}".format(name) not in rpn_left and st.splitter_final:
-                left = ["_{}".format(name), left]
-    else:
-        left = ["_{}".format(name) for name in other_states]
-        if len(left) == 1:
-            left = left[0]
-    return left
-
-
-def _left_right_check(splitter_part, other_states):
-    """
-    Check if splitter_part is purely Left or Right.
-
-    String is returned. If the splitter_part is mixed None is returned.
-
-    """
-    rpn_part = splitter2rpn(
-        splitter_part, other_states=other_states, state_fields=False
-    )
-    inputs_in_splitter = [i for i in rpn_part if i not in ["*", "."]]
-    others_in_splitter = [
-        True if el.startswith("_") else False for el in inputs_in_splitter
-    ]
-    if all(others_in_splitter):
-        return "Left"
-    elif (not all(others_in_splitter)) and (not any(others_in_splitter)):
-        return "Right"
-
-
 def inputs_types_to_dict(name, inputs):
     """Convert type.Inputs to dictionary."""
     # dj: any better option?

pydra/engine/state.py

@@ -97,8 +97,9 @@ def __init__(self, name, splitter=None, combiner=None, other_states=None):
             self.other_states = other_states
         self.splitter = splitter
         # if missing_connections, we can't continue, should wait for updates
+        # TODO: should find a better way, so it's not in the init, but combiner complicates
         if not self.missing_connections:
-            self.connect_splitters()
+            self._connect_splitters()
             self.combiner = combiner
             self.inner_inputs = {}
             for name, (st, inp) in self.other_states.items():
@@ -193,21 +194,15 @@ def combiner(self, combiner):
             self._left_combiner = []
             self._right_combiner = []
 
-    def connect_splitters(self):
+    def _connect_splitters(self):
         """
         Connect splitters from previous nodes.
 
         Evaluates Left (the part from previous states) and Right (current state) parts.
 
         """
         if self.other_states:
-            (
-                self.splitter,
-                self._left_splitter,
-                self._right_splitter,
-            ) = hlpst.connect_splitters(
-                splitter=self.splitter, other_states=self.other_states
-            )
+            self._merge_splitters()
             # left rpn part, but keeping the names of the nodes, e.g. [_NA, _NB, *]
             self._left_splitter_rpn_compact = hlpst.splitter2rpn(
                 deepcopy(self._left_splitter),
@@ -219,13 +214,99 @@ def connect_splitters(self):
             )
         else:  # if other_states is empty there is only Right part
             self._left_splitter = None
-            self._left_splitter_rpn_compact = []
-            self._left_splitter_rpn = []
+            self._left_splitter_rpn_compact, self._left_splitter_rpn = [], []
             self._right_splitter = self.splitter
         self._right_splitter_rpn = hlpst.splitter2rpn(
             deepcopy(self._right_splitter), other_states=self.other_states
         )
 
+    def _merge_splitters(self):
+        """
+        Merging current splitter with the ones from other states.
+
+        If left splitter is not provided the splitter has to be completed.
+
+        """
+        if self.splitter:
+            # if splitter is string, have to check if this is Left or Right part (Left is required)
+            if isinstance(self.splitter, str):
+                # so this is the Left part
+                if self.splitter.startswith("_"):
+                    self._left_splitter = self._complete_left(left=self.splitter)
+                    self._right_splitter = None
+                else:  # this is Right part
+                    self._left_splitter = self._complete_left()
+                    self._right_splitter = self.splitter
+            elif isinstance(self.splitter, (tuple, list)):
+                lr_flag = self._left_right_check(self.splitter)
+                if lr_flag == "Left":
+                    self._left_splitter = self._complete_left(left=self.splitter)
+                    self._right_splitter = None
+                elif lr_flag == "Right":
+                    self._left_splitter = self._complete_left()
+                    self._right_splitter = self.splitter
+                elif lr_flag == "[Left, Right]":
+                    self._left_splitter = self._complete_left(left=self.splitter[0])
+                    self._right_splitter = self.splitter[1]
+        else:
+            # if there is no splitter, I create the Left part
+            self._left_splitter = self._complete_left()
+            self._right_splitter = None
+
+        if self._right_splitter:
+            self.splitter = [
+                deepcopy(self._left_splitter),
+                deepcopy(self._right_splitter),
+            ]
+        else:
+            self.splitter = deepcopy(self._left_splitter)
+
+    def _complete_left(self, left=None):
+        """Add all splitters from previous nodes (completing left part)."""
+        if left:
+            rpn_left = hlpst.splitter2rpn(
+                left, other_states=self.other_states, state_fields=False
+            )
+            for name, (st, inp) in list(self.other_states.items())[::-1]:
+                if "_{}".format(name) not in rpn_left and st.splitter_final:
+                    left = ["_{}".format(name), left]
+        else:
+            left = ["_{}".format(name) for name in self.other_states]
+            if len(left) == 1:
+                left = left[0]
+        return left
+
+    def _left_right_check(self, splitter_part, rec_lev=0):
+        """
+        Check if splitter_part is purely Left, Right
+        or [Left, Right] if the splitter_part is a list (outer splitter)
+
+        String is returned.
+
+        If the splitter_part is mixed exception is raised.
+
+        """
+        rpn_part = hlpst.splitter2rpn(
+            splitter_part, other_states=self.other_states, state_fields=False
+        )
+        inputs_in_splitter = [i for i in rpn_part if i not in ["*", "."]]
+        others_in_splitter = [
+            True if el.startswith("_") else False for el in inputs_in_splitter
+        ]
+        if all(others_in_splitter):
+            return "Left"
+        elif (not all(others_in_splitter)) and (not any(others_in_splitter)):
+            return "Right"
+        elif (
+            isinstance(self.splitter, list)
+            and rec_lev == 0
+            and self._left_right_check(self.splitter[0], rec_lev=1) == "Left"
+            and self._left_right_check(self.splitter[1], rec_lev=1) == "Right"
+        ):
+            return "[Left, Right]"  # Left and Right parts separated in outer scalar
+        else:
+            raise Exception("Left and Right splitters are mixed - splitter invalid")
+
     def set_splitter_final(self):
         """Evaluate a final splitter after combining."""
         _splitter_rpn_final = hlpst.remove_inp_from_splitter_rpn(
@@ -309,7 +390,8 @@ def merge_previous_states(self):
             else:
                 # if no element from st.splitter is in the current combiner,
                 # using st attributes without changes
-                self.keys_final += st.keys_final
+                if st.keys_final:
+                    self.keys_final += st.keys_final
                 group_for_inputs = st.group_for_inputs_final
                 groups_stack = st.groups_stack_final
 

pydra/engine/tests/test_helpers_state.py

@@ -606,90 +606,3 @@ def test_groups_to_input(group_for_inputs, input_for_groups, ndim):
     res = hlpst.converter_groups_to_input(group_for_inputs)
     assert res[0] == input_for_groups
     assert res[1] == ndim
-
-
-@pytest.mark.parametrize(
-    "splitter, other_states, expected_splitter, expected_left, expected_right",
-    [
-        (
-            None,
-            {"NA": (other_states_to_tests(splitter="NA.a"), "b")},
-            "_NA",
-            "_NA",
-            None,
-        ),
-        (
-            "b",
-            {"NA": (other_states_to_tests(splitter="NA.a"), "b")},
-            ["_NA", "b"],
-            "_NA",
-            "b",
-        ),
-        (
-            ("b", "c"),
-            {"NA": (other_states_to_tests(splitter="NA.a"), "b")},
-            ["_NA", ("b", "c")],
-            "_NA",
-            ("b", "c"),
-        ),
-        (
-            None,
-            {
-                "NA": (other_states_to_tests(splitter="NA.a"), "a"),
-                "NB": (other_states_to_tests(splitter="NB.a"), "b"),
-            },
-            ["_NA", "_NB"],
-            ["_NA", "_NB"],
-            None,
-        ),
-        (
-            "b",
-            {
-                "NA": (other_states_to_tests(splitter="NA.a"), "a"),
-                "NB": (other_states_to_tests(splitter="NB.a"), "b"),
-            },
-            [["_NA", "_NB"], "b"],
-            ["_NA", "_NB"],
-            "b",
-        ),
-        (
-            ["_NA", "b"],
-            {
-                "NA": (other_states_to_tests(splitter="NA.a"), "a"),
-                "NB": (other_states_to_tests(splitter="NB.a"), "b"),
-            },
-            [["_NB", "_NA"], "b"],
-            ["_NB", "_NA"],
-            "b",
-        ),
-    ],
-)
-def test_connect_splitters(
-    splitter, other_states, expected_splitter, expected_left, expected_right
-):
-    updated_splitter, left_splitter, right_splitter = hlpst.connect_splitters(
-        splitter, other_states
-    )
-    assert updated_splitter == expected_splitter
-    assert left_splitter == expected_left
-    assert right_splitter == expected_right
-
-
-@pytest.mark.parametrize(
-    "splitter, other_states",
-    [
-        ("_NB", {"NA": (other_states_to_tests(splitter="NA.a"), "b")}),
-        (("_NA", "b"), {"NA": (other_states_to_tests(splitter="NA.a"), "b")}),
-        (["b", "_NA"], {"NA": (other_states_to_tests(splitter="NA.a"), "b")}),
-        (
-            ["_NB", ["_NA", "b"]],
-            {
-                "NA": (other_states_to_tests(splitter="NA.a"), "a"),
-                "NB": (other_states_to_tests(splitter="NB.a"), "b"),
-            },
-        ),
-    ],
-)
-def test_connect_splitters_exception(splitter, other_states):
-    with pytest.raises(Exception):
-        hlpst.connect_splitters(splitter, other_states, state_fields=True)