Add MultiButton support - so you can await on multiple buttons

furbrain · furbrain · commit ae1cedb9affb · 2023-04-10T20:22:47.000+01:00
diff --git a/async_button.py b/async_button.py
@@ -31,11 +31,12 @@
 __repo__ = "https://github.com/furbrain/CircuitPython_async_button.git"
 
 import asyncio
+from asyncio import Task, Event
 
 from adafruit_ticks import ticks_add, ticks_less, ticks_ms
 
 try:
-    from typing import Dict, Sequence, Awaitable
+    from typing import Dict, Sequence, Awaitable, Any, Union
 except ImportError:
     pass
 
@@ -100,6 +101,52 @@ async def released(self):
                 await asyncio.sleep(self.interval)
 
 
+class TaskWrapper:
+    """
+    Create a task to run coro, then trigger event when finished
+    Shouldn't really need this but CircuitPython does not have asyncio.wait or Task.result()
+    """
+
+    def __init__(self, coro: Awaitable, event: Event):
+        """
+        :param Awaitable coro: coroutine to run
+        :param asyncio.Event event:
+        """
+        self._result = None
+        self.coro = coro
+        self.event = event
+        self.task = asyncio.create_task(self._wait())
+
+    async def _wait(self):
+        self._result = await self.coro
+        self.event.set()
+
+    def done(self):
+        """
+        Check whether task has completed
+
+        :return: True if task has completed
+        """
+        return self.task.done()
+
+    def result(self):
+        """
+        Get return value from task
+
+        :return: Whatever the task returned.
+        """
+        if not self.done():
+            raise AssertionError("Task has not yet completed")
+        return self._result
+
+    def cancel(self):
+        """
+        Cancel the task
+        :return:
+        """
+        self.task.cancel()
+
+
 class Button:
     """
     This object will monitor the specified pin for changes and will report
@@ -251,17 +298,12 @@ def _trigger(self, event: int):
         evt.set()
         evt.clear()
 
-    @staticmethod
-    async def _set_event_when_done(coro: Awaitable, event: asyncio.Event):
-        await coro
-        event.set()
-
-    async def wait(self, click_types: Sequence[int] = ALL_EVENTS):
+    async def wait(self, click_types: Union[int, Sequence[int]] = ALL_EVENTS):
         """
         Wait for the first of the specified events.
 
-        :param List[int] click_types: One or more events to listen for. Default is to listen
-          for all events
+        :param (List[int] | int) click_types: List of events to listen for. You can also pass a
+          single event type in. Default is to listen for all events.
         :return: A list of the clicks that actually happened.
 
         :example:
@@ -274,13 +316,19 @@ async def wait(self, click_types: Sequence[int] = ALL_EVENTS):
             >>>         # do something
 
         """
-        evts: Dict[int, asyncio.Task] = {}
+        if isinstance(click_types, int):
+            click_types = [click_types]
+        # shortcut for efficiency: if only one click type awaited, just await that
+        if len(click_types) == 1:
+            await self.events[click_types[0]].wait()
+            return [click_types[0]]
+
+        # multiple click types - needs more complex algorithm
+        evts: Dict[int, TaskWrapper] = {}
         one_event_done = asyncio.Event()
         for evt_type in click_types:
             coro = self.events[evt_type].wait()
-            evts[evt_type] = asyncio.create_task(
-                self._set_event_when_done(coro, one_event_done)
-            )
+            evts[evt_type] = TaskWrapper(coro, one_event_done)
         await one_event_done.wait()
         if len(evts) > 1:
             await asyncio.sleep(0)  # ensure all event types get an opportunity to run
@@ -307,3 +355,58 @@ def deinit(self):
         """
         self.monitor_task.cancel()
         self.keys.deinit()
+
+
+class MultiButton:
+    """
+    This class allows you to await the first click from any of two or more buttons
+    """
+
+    def __init__(self, **kwargs):
+        """
+
+        :param kwargs: pass each button that you want to be able to listen to with its name
+
+        :example:
+          .. code-block:: python
+
+            >>> multi = MultiButton(a = button_a, b=button_b)
+            >>> button, result = await multi.wait(a=Button.SINGLE, b= (Button.DOUBLE, Button.LONG))
+            >>> # Long click on button B
+            >>> print(button, result) # "b", Button.Long
+        """
+        for button in kwargs.values():
+            if not isinstance(button, Button):
+                raise TypeError("Must pass in async_button.Button as parameters")
+        self.buttons: Dict[Any, Button] = kwargs
+
+    async def wait(self, **kwargs):
+        """
+        Wait for any specified clicks
+
+        :param kwargs: pass by keyword what clicks you want to listen for
+        :return: button, click type
+        :example:
+          .. code-block:: python
+
+            >>> multi = MultiButton(a = button_a, b=button_b)
+            >>> button, result = await multi.wait(a=Button.SINGLE, b= (Button.DOUBLE, Button.LONG))
+            >>> # Long click on button B
+            >>> print(button, result) # "b", Button.Long
+        """
+        if len(kwargs) == 1:
+            button = list(kwargs)[0]
+            results = await self.buttons[button].wait(kwargs[button])
+            return button, results[0]
+        tasks: Dict[Any, Task] = {}
+        click_happened = asyncio.Event()
+        for key, value in kwargs.items():
+            tasks[key] = TaskWrapper(self.buttons[key].wait(value), click_happened)
+        await click_happened.wait()
+        result = (None, None)
+        for key, task in tasks.items():
+            if task.done():
+                result = (key, task.result()[0])
+            else:
+                task.cancel()
+        return result
diff --git a/examples/async_multibutton_test.py b/examples/async_multibutton_test.py
@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: 2017 Scott Shawcroft, written for Adafruit Industries
+# SPDX-FileCopyrightText: Copyright (c) 2023 Phil Underwood for Underwood Underground
+#
+# SPDX-License-Identifier: Unlicense
+import asyncio
+
+import board
+
+from async_button import Button, MultiButton
+
+CLICK_NAMES = {
+    Button.SINGLE: "Single click",
+    Button.DOUBLE: "Double click",
+    Button.TRIPLE: "Triple click",
+    Button.LONG: "Long click",
+}
+
+
+async def counter():
+    i = 0
+    while True:
+        print(f"COUNTER: {i}")
+        await asyncio.sleep(1)
+        i += 1
+
+
+async def click_watcher(button: MultiButton):
+    while True:
+        button_name, click = await button.wait(a=Button.ANY_CLICK, b=Button.ANY_CLICK)
+        print(f"{button_name}: {CLICK_NAMES[click]} seen")
+
+
+async def main():
+    # note Button must be created in an async environment
+    button_a = Button(
+        board.D3,
+        value_when_pressed=False,
+        long_click_enable=True,
+    )
+    button_b = Button(
+        board.D4,
+        value_when_pressed=False,
+        long_click_enable=True,
+    )
+    multibutton = MultiButton(a=button_a, b=button_b)
+
+    await asyncio.gather(counter(), click_watcher(multibutton))
+
+
+asyncio.run(main())
diff --git a/tests/test_async_button.py b/tests/test_async_button.py
@@ -120,6 +120,12 @@ async def test_single_click(self):
         await self.wait_event_with_timeout([async_button.Button.SINGLE])
         self.assertAlmostEqual(self.time_count, 0.20, delta=0.1)
 
+    async def test_single_click_specified_without_list(self):
+        self.button = FastButton(self.pin, True)
+        self.button_timings = [0.10, 0.20]
+        await self.wait_event_with_timeout(async_button.Button.SINGLE)
+        self.assertAlmostEqual(self.time_count, 0.20, delta=0.1)
+
     async def test_double_click(self):
         self.button = FastButton(self.pin, True)
         self.button_timings = [0.10, 0.30, 0.5, 0.7]
diff --git a/tests/test_async_multi_button.py b/tests/test_async_multi_button.py
@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 Phil Underwood for Underwood Underground
+#
+# SPDX-License-Identifier: MIT
+from functools import partial
+from unittest import IsolatedAsyncioTestCase
+from unittest.mock import MagicMock, AsyncMock
+import sys
+import asyncio
+
+sys.modules["countio"] = MagicMock()
+
+import async_button  # pylint: disable=wrong-import-position
+
+SINGLE = async_button.Button.SINGLE
+DOUBLE = async_button.Button.DOUBLE
+LONG = async_button.Button.LONG
+
+
+async def wait_and_return(delay: float, click_type):
+    await asyncio.sleep(delay)
+    if isinstance(click_type, int):
+        click_type = [click_type]
+    return click_type
+
+
+class TestButton(IsolatedAsyncioTestCase):
+    # pylint: disable=invalid-name, too-many-public-methods
+    def setUp(self) -> None:
+        self.button_a = MagicMock(async_button.Button)
+        self.button_a.wait = AsyncMock()
+        self.button_b = MagicMock(async_button.Button)
+        self.button_b.wait = AsyncMock()
+        self.button_c = MagicMock(async_button.Button)
+        self.button_c.wait = AsyncMock()
+
+    def testInitialise(self):
+        async_button.MultiButton(a=self.button_a, b=self.button_b, c=self.button_c)
+
+    def testInitialiseFailsWithNonButtons(self):
+        with self.assertRaises(TypeError):
+            async_button.MultiButton(a=self.button_a, b=self.button_b, c=12)
+
+    async def testSimpleCase(self):
+        multi = async_button.MultiButton(a=self.button_a)
+        self.button_a.wait.return_value = [SINGLE]
+        result = await multi.wait(a=SINGLE)
+        self.assertEqual(("a", SINGLE), result)
+        self.button_a.wait.assert_awaited_once()
+        self.button_a.wait.assert_called_with(SINGLE)
+
+    async def testTwoButtons(self):
+        multi = async_button.MultiButton(a=self.button_a, b=self.button_b)
+        self.button_a.wait = partial(wait_and_return, 0.2)
+        self.button_b.wait = partial(wait_and_return, 0.1)
+        result = await multi.wait(a=SINGLE, b=DOUBLE)
+        self.assertEqual(("b", DOUBLE), result)
+
+    async def testThreeButtons(self):
+        multi = async_button.MultiButton(
+            a=self.button_a, b=self.button_b, c=self.button_c
+        )
+        self.button_a.wait = partial(wait_and_return, 0.2)
+        self.button_b.wait = partial(wait_and_return, 0.1)
+        self.button_c.wait = partial(wait_and_return, 0.3)
+        result = await multi.wait(a=SINGLE, b=DOUBLE, c=LONG)
+        self.assertEqual(("b", DOUBLE), result)
+
+    async def testOneButtonTwoClicks(self):
+        multi = async_button.MultiButton(a=self.button_a)
+        self.button_a.wait.return_value = [DOUBLE]
+        result = await multi.wait(a=[SINGLE, DOUBLE])
+        self.assertEqual(("a", DOUBLE), result)
+        self.button_a.wait.assert_awaited_once()
+        self.button_a.wait.assert_called_with([SINGLE, DOUBLE])
