Add ForEach decorator to iterate over a list (#476)

Author: corot

py_trees/decorators.py

@@ -36,6 +36,7 @@
 * :class:`py_trees.decorators.Condition`
 * :class:`py_trees.decorators.Count`
 * :class:`py_trees.decorators.EternalGuard`
+* :class:`py_trees.decorators.ForEach`
 * :class:`py_trees.decorators.Inverter`
 * :class:`py_trees.decorators.OneShot`
 * :class:`py_trees.decorators.Repeat`
@@ -80,6 +81,7 @@
 import inspect
 import time
 import typing
+from collections.abc import Iterable
 
 from . import behaviour, blackboard, common
 
@@ -920,3 +922,74 @@ def update(self) -> common.Status:
             the behaviour's new status :class:`~py_trees.common.Status`
         """
         return self.decorated.status
+
+
+class ForEach(Decorator):
+    """
+    Run the child behavior for each item in an iterable.
+
+    On initialization, the iterable is loaded from the blackboard and an
+    iterator is created. Every time the child succeeds, we advance to
+    the next item. We keep running until the iterable is exhausted.
+    """
+
+    def __init__(
+        self, name: str, child: behaviour.Behaviour, source_key: str, target_key: str
+    ):
+        """
+        Initialise the ForEach decorator.
+
+        Args:
+            name (:obj:`str`): name of the behaviour
+            child (:obj:`Behaviour`): the child behaviour to decorate
+            source_key (:obj:`str`): blackboard key to read the input iterable
+            target_key (:obj:`str`): blackboard key to set for each iteration
+        """
+        super().__init__(name=name, child=child)
+        self.source_key = source_key
+        self.target_key = target_key
+        self.blackboard = blackboard.Client(name=name)
+        self.blackboard.register_key(key=self.source_key, access=common.Access.READ)
+        self.blackboard.register_key(key=self.target_key, access=common.Access.WRITE)
+        self._iterator: typing.Iterator | None = None
+        self._current_item: typing.Any | None = None
+
+    def initialise(self) -> None:
+        """Reset iteration on first tick."""
+        iterable = self.blackboard.get(self.source_key) or []
+        if not isinstance(iterable, Iterable):
+            raise TypeError(
+                f"[{self.name}] source_key '{self.source_key}' is not an iterable"
+            )
+        self._iterator = iter(iterable)
+        self._advance()
+
+    def _advance(self) -> None:
+        """Advance to the next item in the iterator, or mark as finished."""
+        try:
+            if self._iterator is not None:
+                self._current_item = next(self._iterator)
+                self.blackboard.set(self.target_key, self._current_item)
+        except StopIteration:
+            self._current_item = None
+
+    def update(self) -> common.Status:
+        """Execute the child for the current item and manage iteration state."""
+        if self._current_item is None:
+            # no more items left
+            return common.Status.SUCCESS
+
+        child_status = self.decorated.status
+
+        if child_status == common.Status.SUCCESS:
+            # move to next item
+            self._advance()
+            if self._current_item is not None:
+                return common.Status.RUNNING
+
+        return child_status
+
+    def terminate(self, new_status: common.Status) -> None:
+        """Reset on termination."""
+        self._iterator = None
+        self._current_item = None

tests/test_decorators.py

@@ -9,6 +9,7 @@
 ##############################################################################
 
 import time
+import typing
 
 import py_trees
 import py_trees.console as console
@@ -604,3 +605,159 @@ def test_status_to_blackboard() -> None:
         result=decorator.status,
     )
     assert decorator.status == py_trees.common.Status.SUCCESS
+
+
+@pytest.mark.parametrize(
+    "iterable",
+    [
+        [1, 2, 3],  # list
+        (1, 2, 3),  # tuple
+        (i for i in range(1, 4)),  # generator
+    ],
+    ids=["list", "tuple", "generator"],
+)
+def test_for_each_fixed(iterable: typing.Iterable[int]) -> None:
+    console.banner("ForEach - 3 elements on multiple iterables")
+    child = py_trees.behaviours.StatusQueue(
+        name="R-S",
+        queue=[
+            py_trees.common.Status.RUNNING,
+            py_trees.common.Status.SUCCESS,
+        ],
+        eventually=None,
+    )
+    blackboard = py_trees.blackboard.Client()
+    blackboard.register_key(key="element", access=py_trees.common.Access.READ)
+    blackboard.register_key(key="iterable", access=py_trees.common.Access.WRITE)
+    blackboard.iterable = iterable
+    decorator = py_trees.decorators.ForEach(
+        name="ForEach", child=child, source_key="iterable", target_key="element"
+    )
+
+    decorator.tick_once()  # run first element
+    print("\n--------- Tick 1 ---------\n")
+    print("decorator.status == py_trees.common.Status.RUNNING")
+    assert decorator.status == py_trees.common.Status.RUNNING
+    print("child.status == py_trees.common.Status.RUNNING")
+    assert child.status == py_trees.common.Status.RUNNING
+    assert blackboard.element == 1
+
+    decorator.tick_once()  # complete first element; load second
+    print("\n--------- Tick 2 ---------\n")
+    print("decorator.status == py_trees.common.Status.RUNNING")
+    assert decorator.status == py_trees.common.Status.RUNNING
+    print("child.status == py_trees.common.Status.SUCCESS")
+    assert child.status == py_trees.common.Status.SUCCESS
+    assert blackboard.element == 2  # type: ignore[unreachable]
+
+    decorator.tick_once()  # run second element
+    print("\n--------- Tick 3 ---------\n")
+    print("decorator.status == py_trees.common.Status.RUNNING")
+    assert decorator.status == py_trees.common.Status.RUNNING
+    print("child.status == py_trees.common.Status.RUNNING")
+    assert child.status == py_trees.common.Status.RUNNING
+    assert blackboard.element == 2
+
+    decorator.tick_once()  # complete second element; load third
+    print("\n--------- Tick 4 ---------\n")
+    print("decorator.status == py_trees.common.Status.RUNNING")
+    assert decorator.status == py_trees.common.Status.RUNNING
+    print("child.status == py_trees.common.Status.SUCCESS")
+    assert child.status == py_trees.common.Status.SUCCESS
+    assert blackboard.element == 3
+
+    decorator.tick_once()  # run third element
+    print("\n--------- Tick 5 ---------\n")
+    print("decorator.status == py_trees.common.Status.RUNNING")
+    assert decorator.status == py_trees.common.Status.RUNNING
+    print("child.status == py_trees.common.Status.RUNNING")
+    assert child.status == py_trees.common.Status.RUNNING
+    assert blackboard.element == 3
+
+    decorator.tick_once()  # complete third element
+    print("\n--------- Tick 6 ---------\n")
+    print("decorator.status == py_trees.common.Status.SUCCESS")
+    assert decorator.status == py_trees.common.Status.SUCCESS
+    print("child.status == py_trees.common.Status.SUCCESS")
+    assert child.status == py_trees.common.Status.SUCCESS
+    assert blackboard.element == 3
+
+
+def test_for_each_append() -> None:
+    console.banner("ForEach - 3 + 2 elements")
+    child = py_trees.behaviours.StatusQueue(
+        name="S",
+        queue=[
+            py_trees.common.Status.SUCCESS,
+        ],
+        eventually=None,
+    )
+    blackboard = py_trees.blackboard.Client()
+    blackboard.register_key(key="element", access=py_trees.common.Access.READ)
+    blackboard.register_key(key="iterable", access=py_trees.common.Access.WRITE)
+    blackboard.iterable = [1, 2, 3]
+    decorator = py_trees.decorators.ForEach(
+        name="ForEach", child=child, source_key="iterable", target_key="element"
+    )
+
+    decorator.tick_once()
+    print("\n--------- Tick 1 ---------\n")
+    assert blackboard.element == 2
+
+    decorator.tick_once()
+    print("\n--------- Tick 2 ---------\n")
+    assert blackboard.element == 3
+
+    # add two elements, so we need two more ticks to succeed
+    blackboard.iterable.extend([4, 5])
+
+    decorator.tick_once()
+    print("\n--------- Tick 3 ---------\n")
+    assert blackboard.element == 4
+
+    decorator.tick_once()
+    print("\n--------- Tick 4 ---------\n")
+    assert blackboard.element == 5
+
+    decorator.tick_once()
+    print("\n--------- Tick 5 ---------\n")
+    print("decorator.status == py_trees.common.Status.SUCCESS")
+    assert decorator.status == py_trees.common.Status.SUCCESS
+    print("child.status == py_trees.common.Status.SUCCESS")
+    assert child.status == py_trees.common.Status.SUCCESS
+
+
+def test_for_each_delete() -> None:
+    console.banner("ForEach - 4 + 1 elements")
+    child = py_trees.behaviours.StatusQueue(
+        name="S",
+        queue=[
+            py_trees.common.Status.SUCCESS,
+        ],
+        eventually=None,
+    )
+    blackboard = py_trees.blackboard.Client()
+    blackboard.register_key(key="element", access=py_trees.common.Access.READ)
+    blackboard.register_key(key="iterable", access=py_trees.common.Access.WRITE)
+    blackboard.iterable = [1, 2, 3, 4]
+    decorator = py_trees.decorators.ForEach(
+        name="ForEach", child=child, source_key="iterable", target_key="element"
+    )
+
+    decorator.tick_once()
+    print("\n--------- Tick 1 ---------\n")
+    assert blackboard.element == 2
+
+    decorator.tick_once()
+    print("\n--------- Tick 2 ---------\n")
+    assert blackboard.element == 3
+
+    # remove last element, so we skip one tick to succeed
+    blackboard.iterable.pop()
+
+    decorator.tick_once()
+    print("\n--------- Tick 3 ---------\n")
+    print("decorator.status == py_trees.common.Status.SUCCESS")
+    assert decorator.status == py_trees.common.Status.SUCCESS
+    print("child.status == py_trees.common.Status.SUCCESS")
+    assert child.status == py_trees.common.Status.SUCCESS