Repository: bluesky/bluesky
Pull request #1804 (2024-09-28 11:33:09)

Windows: allow use of `bps.input_plan()`
## Description

Implement `AsyncInput` using `run_in_executor`, as `loop.add_reader()` is not available for `stdin` on either of the standard windows event loop implementations.

`run_in_executor` may add a small amount of overhead (it uses a thread pool under the hood), but since we're then immediately asking for user input I do not think that is significant (if other facilities feel it _is_ a problem, we can split the implementation into windows-specific and linux-specific versions, but I don't feel it's worth the complexity cost personally...)

## Motivation and Context

At ISIS we're using bluesky on windows. On windows, `bps.input_plan()` doesn't currently work with either the proactor or selector event loop implementations. See https://docs.python.org/3/library/asyncio-platforms.html#asyncio-platform-support for details about which asyncio methods are different/unavailable on windows.

```
Python 3.11.6 (tags/v3.11.6:8b6ee5b, Oct  2 2023, 14:57:12) [MSC v.1935 64 bit (AMD64)] on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> from bluesky.run_engine import RunEngine
>>> import bluesky.plan_stubs as bps
>>> RE = RunEngine()
>>> RE(bps.input_plan("foo?"))
Run aborted
Traceback (most recent call last):
  File "C:\Instrument\dev\bluesky\src\bluesky\run_engine.py", line 1580, in _run
    msg = self._plan_stack[-1].throw(stashed_exception or resp)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Instrument\dev\bluesky\src\bluesky\utils\__init__.py", line 1949, in __iter__
    return (yield from self._iter)
            ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Instrument\dev\bluesky\src\bluesky\plan_stubs.py", line 745, in input_plan
    return (yield Msg("input", prompt=prompt))
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Instrument\dev\bluesky\src\bluesky\run_engine.py", line 1672, in _run
    new_response = await coro(msg)
                   ^^^^^^^^^^^^^^^
  File "C:\Instrument\dev\bluesky\src\bluesky\run_engine.py", line 2668, in _input
    async_input = AsyncInput(self.loop)
                  ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Instrument\dev\bluesky\src\bluesky\utils\__init__.py", line 1165, in __init__
    self.loop.add_reader(sys.stdin, self.got_input)
  File "c:\instrument\apps\python3\Lib\asyncio\events.py", line 530, in add_reader
    raise NotImplementedError
NotImplementedError
```

## How Has This Been Tested?

Unit testing this is a bit gnarly, but have had a go at implementing unit tests for both the high level `input_plan`, and checking that the low-level `AsyncInput` object works asynchronously and doesn't block the event loop. Suggestions to make the unit tests a bit nicer appreciated...

Have also manually tested `input_plan()` on both windows & linux.

Doesn't play well with ctrl-c, but it's no _worse_ than what was previously on `main` in that regard. May have a look at this in a follow-up PR.


Patch:
diff --git a/src/bluesky/utils/__init__.py b/src/bluesky/utils/__init__.py
index c7c839b6..466352ea 100644
--- a/src/bluesky/utils/__init__.py
+++ b/src/bluesky/utils/__init__.py
@@ -1151,25 +1151,12 @@ def register_transform(RE, *, prefix="<", ip=None):
 
 
 class AsyncInput:
-    """a input prompt that allows event loop to run in the background
-
-    adapted from http://stackoverflow.com/a/35514777/1221924
-    """
-
     def __init__(self, loop=None):
         self.loop = loop or asyncio.get_event_loop()
-        if sys.version_info < (3, 10):
-            self.q = asyncio.Queue(loop=self.loop)
-        else:
-            self.q = asyncio.Queue()
-        self.loop.add_reader(sys.stdin, self.got_input)
-
-    def got_input(self):
-        asyncio.ensure_future(self.q.put(sys.stdin.readline()), loop=self.loop)
 
     async def __call__(self, prompt, end="\n", flush=False):
         print(prompt, end=end, flush=flush)
-        return (await self.q.get()).rstrip("\n")
+        return (await self.loop.run_in_executor(None, sys.stdin.readline)).rstrip("\n")
 
 
 def new_uid():


Test patch:
diff --git a/src/bluesky/tests/test_plans.py b/src/bluesky/tests/test_plans.py
index 0eccba0f..bf3f55b7 100644
--- a/src/bluesky/tests/test_plans.py
+++ b/src/bluesky/tests/test_plans.py
@@ -2,6 +2,7 @@ import collections
 import inspect
 import random
 import re
+from unittest.mock import patch
 
 import numpy as np
 import numpy.testing as npt
@@ -712,3 +713,12 @@ def test_count_failure(RE, hw):
 def test_no_warning(hw):
     # this should not warn
     list(bps.rel_set(hw.motor, 2))
+
+
+def test_input_plan(RE):
+    def plan():
+        result = yield from bps.input_plan("prompt: ")
+        assert result == "answer"
+
+    with patch("bluesky.utils.sys.stdin.readline", return_value="answer\n"):
+        RE(plan())
diff --git a/src/bluesky/tests/test_utils.py b/src/bluesky/tests/test_utils.py
index c8f4a18c..6095a959 100644
--- a/src/bluesky/tests/test_utils.py
+++ b/src/bluesky/tests/test_utils.py
@@ -1,6 +1,9 @@
+import asyncio
 import operator
+import time
 import warnings
 from functools import reduce
+from unittest.mock import patch
 
 import numpy as np
 import pytest
@@ -8,7 +11,9 @@ from cycler import cycler
 
 from bluesky import RunEngine
 from bluesky.plan_stubs import complete_all, mv
+from bluesky.run_engine import WaitForTimeoutError
 from bluesky.utils import (
+    AsyncInput,
     CallbackRegistry,
     Msg,
     ensure_generator,
@@ -552,3 +557,45 @@ def test_warning_behavior(gen_func, iterated):
 )
 def test_check_if_func_is_plan(func, is_plan_result):
     assert is_plan(func) == is_plan_result
+
+
+def test_async_input_does_not_block_event_loop(RE, capsys):
+    a_short_time = 0.5
+    background_event = asyncio.Event()
+    input_completed_event = asyncio.Event()
+
+    async def get_input():
+        ai = AsyncInput()
+
+        # Patch stdin.readline as a *blocking* function that takes a while to return - the wait_for
+        # should time out before this returns, so the input_completed_event should never have a
+        # chance to be set.
+        with patch("bluesky.utils.sys.stdin.readline", side_effect=lambda: time.sleep(3 * a_short_time)):
+            await ai("prompt: ")
+
+        input_completed_event.set()
+
+    async def sleep_then_set_background_event():
+        await asyncio.sleep(a_short_time)
+        background_event.set()
+
+    def plan():
+        get_input_fut = asyncio.ensure_future(get_input(), loop=RE.loop)
+        try:
+            yield Msg(
+                "wait_for",
+                None,
+                [lambda: get_input_fut, lambda: sleep_then_set_background_event()],
+                timeout=a_short_time * 2,
+            )
+        except WaitForTimeoutError:
+            # Expected
+            # Input call should not have managed to "complete" before timeout in test(), but the
+            # background event should have been set as the input call shouldn't stall the
+            # whole event loop.
+            assert background_event.is_set() and not input_completed_event.is_set()
+            get_input_fut.cancel()
+        else:
+            pytest.fail("Should have timed out waiting for input")
+
+    RE(plan())