Merge pull request #671 from isidentical/fsspec-loop-ctx

 asyn: introduce fsspec.asyn.fsspec_loop

Author: martindurant

docs/source/async.rst

@@ -127,3 +127,5 @@ available as the attribute ``.loop``.
 .. autofunction:: fsspec.asyn.sync_wrapper
 
 .. autofunction:: fsspec.asyn.get_loop
+
+.. autofunction:: fsspec.asyn.fsspec_loop

docs/source/changelog.rst

@@ -4,7 +4,9 @@ Changelog
 Dev
 ---
 
+Enhancements
 
+- Introduce ``fsspec.asyn.fsspec_loop`` to temporarily switch to the fsspec loop. (#671)
 
 
 2021.06.0

fsspec/asyn.py

@@ -105,6 +105,17 @@ def _selector_policy():
         asyncio.set_event_loop_policy(original_policy)
 
 
+def get_running_loop():
+    if hasattr(asyncio, "get_running_loop"):
+        return asyncio.get_running_loop()
+    else:
+        loop = asyncio._get_running_loop()
+        if loop is None:
+            raise RuntimeError("no running event loop")
+        else:
+            return loop
+
+
 def get_loop():
     """Create or return the default fsspec IO loop
 
@@ -124,6 +135,25 @@ def get_loop():
     return loop[0]
 
 
+@contextmanager
+def fsspec_loop():
+    """Temporarily switch the current event loop to the fsspec's
+    own loop, and then revert it back after the context gets
+    terinated.
+    """
+    try:
+        original_loop = get_running_loop()
+    except RuntimeError:
+        original_loop = None
+
+    fsspec_loop = get_loop()
+    try:
+        asyncio._set_running_loop(fsspec_loop)
+        yield fsspec_loop
+    finally:
+        asyncio._set_running_loop(original_loop)
+
+
 try:
     import resource
 except ImportError:

fsspec/tests/test_async.py

@@ -8,7 +8,7 @@
 
 import fsspec
 import fsspec.asyn
-from fsspec.asyn import _throttled_gather
+from fsspec.asyn import _throttled_gather, get_running_loop
 
 
 def test_sync_methods():
@@ -132,3 +132,23 @@ def test_windows_policy():
     # check ensures that we are restoring the old policy back
     # after our change.
     assert isinstance(policy, asyncio.DefaultEventLoopPolicy)
+
+
+def test_fsspec_loop():
+    asyncio._set_running_loop(None)
+
+    with fsspec.asyn.fsspec_loop() as loop:
+        assert get_running_loop() is loop
+        assert get_running_loop() is fsspec.asyn.get_loop()
+
+    with pytest.raises(RuntimeError):
+        get_running_loop()
+
+    original_loop = asyncio.new_event_loop()
+    asyncio._set_running_loop(original_loop)
+
+    with fsspec.asyn.fsspec_loop() as loop:
+        assert get_running_loop() is loop
+        assert get_running_loop() is fsspec.asyn.get_loop()
+
+    assert get_running_loop() is original_loop