Make get_event_loop() return the current loop if called from coroutines

Author: 1st1

asyncio/base_events.py

@@ -400,6 +400,8 @@ def run_forever(self):
             old_agen_hooks = sys.get_asyncgen_hooks()
             sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter_hook,
                                    finalizer=self._asyncgen_finalizer_hook)
+        old_loop = events._get_current_loop()
+        events._set_current_loop(self)
         try:
             while True:
                 self._run_once()
@@ -408,6 +410,7 @@ def run_forever(self):
         finally:
             self._stopping = False
             self._thread_id = None
+            events._set_current_loop(old_loop)
             self._set_coroutine_wrapper(False)
             if self._asyncgens is not None:
                 sys.set_asyncgen_hooks(*old_agen_hooks)

asyncio/events.py

@@ -607,6 +607,27 @@ def new_event_loop(self):
 # Lock for protecting the on-the-fly creation of the event loop policy.
 _lock = threading.Lock()
 
+# A TLS for the current event loop, used by get_current_loop.
+_current_loop = threading.local()
+
+
+def _get_current_loop():
+    """Return the currently running event loop or None.
+
+    This is a low-level function intended to be used by event loops.
+    This function is thread-specific.
+    """
+    return getattr(_current_loop, 'loop', None)
+
+
+def _set_current_loop(loop):
+    """Set the currently running event loop.
+
+    This is a low-level function intended to be used by event loops.
+    This function is thread-specific.
+    """
+    _current_loop.loop = loop
+
 
 def _init_event_loop_policy():
     global _event_loop_policy
@@ -633,7 +654,17 @@ def set_event_loop_policy(policy):
 
 
 def get_event_loop():
-    """Equivalent to calling get_event_loop_policy().get_event_loop()."""
+    """Return an asyncio event loop.
+
+    When called from coroutines, this function will always return the
+    current event loop.
+
+    If there is no current event loop set, the function will return
+    the result of `get_event_loop_policy().get_event_loop()` call.
+    """
+    current_loop = _get_current_loop()
+    if current_loop is not None:
+        return current_loop
     return get_event_loop_policy().get_event_loop()
 
 

asyncio/test_utils.py

@@ -449,7 +449,13 @@ def new_test_loop(self, gen=None):
         self.set_event_loop(loop)
         return loop
 
+    def setUp(self):
+        self._get_current_loop = events._get_current_loop
+        events._get_current_loop = lambda: None
+
     def tearDown(self):
+        events._get_current_loop = self._get_current_loop
+
         events.set_event_loop(None)
 
         # Detect CPython bug #23353: ensure that yield/yield-from is not used

tests/test_base_events.py

@@ -154,6 +154,7 @@ def test_ipaddr_info_no_inet_pton(self, m_socket):
 class BaseEventLoopTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = base_events.BaseEventLoop()
         self.loop._selector = mock.Mock()
         self.loop._selector.select.return_value = ()
@@ -997,6 +998,7 @@ def connection_lost(self, exc):
 class BaseEventLoopWithSelectorTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = asyncio.new_event_loop()
         self.set_event_loop(self.loop)
 

tests/test_events.py

@@ -2233,6 +2233,7 @@ def noop(*args, **kwargs):
 class HandleTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = mock.Mock()
         self.loop.get_debug.return_value = True
 
@@ -2418,6 +2419,7 @@ def __await__(self):
 class TimerTests(unittest.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = mock.Mock()
 
     def test_hash(self):
@@ -2726,6 +2728,27 @@ def test_set_event_loop_policy(self):
         self.assertIs(policy, asyncio.get_event_loop_policy())
         self.assertIsNot(policy, old_policy)
 
+    def test_get_event_loop_returns_current_loop(self):
+        class Policy(asyncio.DefaultEventLoopPolicy):
+            def get_event_loop(self):
+                raise NotImplementedError
+
+        loop = None
+
+        old_policy = asyncio.get_event_loop_policy()
+        asyncio.set_event_loop_policy(Policy())
+        try:
+            loop = asyncio.new_event_loop()
+
+            async def func():
+                self.assertIs(asyncio.get_event_loop(), loop)
+
+            loop.run_until_complete(func())
+        finally:
+            asyncio.set_event_loop_policy(old_policy)
+            if loop is not None:
+                loop.close()
+
 
 if __name__ == '__main__':
     unittest.main()

tests/test_futures.py

@@ -79,6 +79,7 @@ def __iter__(self):
 class DuckTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = self.new_test_loop()
         self.addCleanup(self.loop.close)
 
@@ -96,6 +97,7 @@ def test_ensure_future(self):
 class FutureTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = self.new_test_loop()
         self.addCleanup(self.loop.close)
 
@@ -468,6 +470,7 @@ def test_set_result_unless_cancelled(self):
 class FutureDoneCallbackTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = self.new_test_loop()
 
     def run_briefly(self):

tests/test_locks.py

@@ -19,6 +19,7 @@
 class LockTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = self.new_test_loop()
 
     def test_ctor_loop(self):
@@ -235,6 +236,7 @@ def test_context_manager_no_yield(self):
 class EventTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = self.new_test_loop()
 
     def test_ctor_loop(self):
@@ -364,6 +366,7 @@ def c1(result):
 class ConditionTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = self.new_test_loop()
 
     def test_ctor_loop(self):
@@ -699,6 +702,7 @@ def test_ambiguous_loops(self):
 class SemaphoreTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = self.new_test_loop()
 
     def test_ctor_loop(self):

tests/test_pep492.py

@@ -17,6 +17,7 @@
 class BaseTest(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = asyncio.BaseEventLoop()
         self.loop._process_events = mock.Mock()
         self.loop._selector = mock.Mock()

tests/test_proactor_events.py

@@ -24,6 +24,7 @@ def close_transport(transport):
 class ProactorSocketTransportTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = self.new_test_loop()
         self.addCleanup(self.loop.close)
         self.proactor = mock.Mock()
@@ -436,6 +437,8 @@ def test_dont_pause_writing(self):
 class BaseProactorEventLoopTests(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
+
         self.sock = test_utils.mock_nonblocking_socket()
         self.proactor = mock.Mock()
 

tests/test_queues.py

@@ -10,6 +10,7 @@
 class _QueueTestBase(test_utils.TestCase):
 
     def setUp(self):
+        super().setUp()
         self.loop = self.new_test_loop()