Add AsyncRunner.block_until_done (#444)

basnijholt · pre-commit-ci[bot] · web-flow · commit bdcbbf35d808 · 2024-02-13T10:59:49.000-08:00
* Add `AsyncRunner.block_until_done` I have answered this question so often that I think it warrents a new method. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use runner.block_until_done --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
diff --git a/adaptive/runner.py b/adaptive/runner.py
@@ -722,6 +722,14 @@ def cancel(self) -> None:
         """
         self.task.cancel()
 
+    def block_until_done(self) -> None:
+        if in_ipynb():
+            raise RuntimeError(
+                "Cannot block the event loop when running in a Jupyter notebook."
+                " Use `await runner.task` instead."
+            )
+        self.ioloop.run_until_complete(self.task)
+
     def live_plot(
         self,
         *,
diff --git a/adaptive/tests/test_runner.py b/adaptive/tests/test_runner.py
@@ -1,4 +1,3 @@
-import asyncio
 import platform
 import sys
 import time
@@ -34,7 +33,7 @@ def blocking_runner(learner, **kw):
 
 def async_runner(learner, **kw):
     runner = AsyncRunner(learner, executor=SequentialExecutor(), **kw)
-    asyncio.get_event_loop().run_until_complete(runner.task)
+    runner.block_until_done()
 
 
 runners = [simple, blocking_runner, async_runner]
@@ -71,7 +70,7 @@ async def f(x):
 
     learner = Learner1D(f, (-1, 1))
     runner = AsyncRunner(learner, npoints_goal=10)
-    asyncio.get_event_loop().run_until_complete(runner.task)
+    runner.block_until_done()
 
 
 # --- Test with different executors
@@ -158,7 +157,7 @@ def test_loky_executor(loky_executor):
 def test_default_executor():
     learner = Learner1D(linear, (-1, 1))
     runner = AsyncRunner(learner, npoints_goal=10)
-    asyncio.get_event_loop().run_until_complete(runner.task)
+    runner.block_until_done()
 
 
 def test_auto_goal():
diff --git a/docs/source/tutorial/tutorial.parallelism.md b/docs/source/tutorial/tutorial.parallelism.md
@@ -89,7 +89,7 @@ if __name__ == "__main__":
     runner.start_periodic_saving(dict(fname=fname), interval=600)
 
     # block until runner goal reached
-    runner.ioloop.run_until_complete(runner.task)
+    runner.block_until_done()
 
     # save one final time before exiting
     learner.save(fname)
