Merge pull request #13349 from minrk/policy_get_event_loop

Carreau · web-flow · commit e343cd4e5872 · 2021-12-01T11:38:43.000-08:00
get_running_loop is only valid in coroutines
diff --git a/IPython/core/async_helpers.py b/IPython/core/async_helpers.py
@@ -12,17 +12,28 @@
 
 
 import ast
+import asyncio
 import inspect
 
 
 class _AsyncIORunner:
+    def __init__(self):
+        self._loop = None
+
+    @property
+    def loop(self):
+        """Always returns a non-closed event loop"""
+        if self._loop is None or self._loop.is_closed():
+            policy = asyncio.get_event_loop_policy()
+            self._loop = policy.new_event_loop()
+            policy.set_event_loop(self._loop)
+        return self._loop
+
     def __call__(self, coro):
         """
         Handler for asyncio autoawait
         """
-        import asyncio
-
-        return asyncio.get_event_loop_policy().get_event_loop().run_until_complete(coro)
+        return self.loop.run_until_complete(coro)
 
     def __str__(self):
         return "asyncio"
diff --git a/IPython/core/magics/script.py b/IPython/core/magics/script.py
@@ -81,7 +81,15 @@ def safe_watcher():
         yield
         return
 
-    loop = policy.get_event_loop()
+    try:
+        loop = policy.get_event_loop()
+        if loop.is_closed():
+            raise RuntimeError("open a new one")
+    except RuntimeError:
+        # closed loop, make a new one
+        loop = policy.new_event_loop()
+        policy.set_event_loop(loop)
+
     try:
         watcher = asyncio.SafeChildWatcher()
         watcher.attach_loop(loop)
@@ -236,9 +244,19 @@ async def _stream_communicate(process, cell):
             await asyncio.wait([stdout_task, stderr_task])
             await process.wait()
 
-        if sys.platform.startswith("win"):
-            asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
-        loop = asyncio.get_event_loop_policy().get_event_loop()
+        policy = asyncio.get_event_loop_policy()
+        if sys.platform.startswith("win") and not isinstance(
+            policy, asyncio.WindowsProactorEventLoopPolicy
+        ):
+            # _do not_ overwrite the current policy
+            policy = asyncio.WindowsProactorEventLoopPolicy()
+
+        try:
+            loop = policy.get_event_loop()
+        except RuntimeError:
+            # closed loop, make a new one
+            loop = policy.new_event_loop()
+            policy.set_event_loop(loop)
         argv = arg_split(line, posix=not sys.platform.startswith("win"))
         args, cmd = self.shebang.parser.parse_known_args(argv)
         try:
diff --git a/IPython/core/tests/test_interactiveshell.py b/IPython/core/tests/test_interactiveshell.py
@@ -1058,6 +1058,22 @@ def test_run_cell_async():
     assert result.result == 5
 
 
+def test_run_cell_await():
+    ip.run_cell("import asyncio")
+    result = ip.run_cell("await asyncio.sleep(0.01); 10")
+    assert ip.user_ns["_"] == 10
+
+
+def test_run_cell_asyncio_run():
+    ip.run_cell("import asyncio")
+    result = ip.run_cell("await asyncio.sleep(0.01); 1")
+    assert ip.user_ns["_"] == 1
+    result = ip.run_cell("asyncio.run(asyncio.sleep(0.01)); 2")
+    assert ip.user_ns["_"] == 2
+    result = ip.run_cell("await asyncio.sleep(0.01); 3")
+    assert ip.user_ns["_"] == 3
+
+
 def test_should_run_async():
     assert not ip.should_run_async("a = 5")
     assert ip.should_run_async("await x")
diff --git a/IPython/core/tests/test_magic.py b/IPython/core/tests/test_magic.py
@@ -970,7 +970,11 @@ def test_script_config():
 
 @pytest.fixture
 def event_loop():
-    yield asyncio.get_event_loop_policy().get_event_loop()
+    policy = asyncio.get_event_loop_policy()
+    loop = policy.new_event_loop()
+    policy.set_event_loop(loop)
+    yield loop
+    loop.close()
 
 
 @dec.skip_win32
diff --git a/IPython/terminal/interactiveshell.py b/IPython/terminal/interactiveshell.py
@@ -504,21 +504,24 @@ def prompt_for_code(self):
         # If we don't do this, people could spawn coroutine with a
         # while/true inside which will freeze the prompt.
 
+        policy = asyncio.get_event_loop_policy()
         try:
-            old_loop = asyncio.get_running_loop()
+            old_loop = policy.get_event_loop()
         except RuntimeError:
-            # This happens when the user used `asyncio.run()`.
+            # This happens when the the event loop is closed,
+            # e.g. by calling `asyncio.run()`.
             old_loop = None
 
-        asyncio.set_event_loop(self.pt_loop)
+        policy.set_event_loop(self.pt_loop)
         try:
             with patch_stdout(raw=True):
                 text = self.pt_app.prompt(
                     default=default,
                     **self._extra_prompt_options())
         finally:
             # Restore the original event loop.
-            asyncio.set_event_loop(old_loop)
+            if old_loop is not None:
+                policy.set_event_loop(old_loop)
 
         return text
 
