feat: asyncio LoopExecutor and async fsspec source

src/uproot/source/fsspec.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import concurrent.futures
 import queue
 
@@ -23,6 +24,7 @@ class FSSpecSource(uproot.source.chunk.Source):
     """
 
     def __init__(self, file_path: str, **options):
+        import fsspec.asyn
         import fsspec.core
 
         default_options = uproot.reading.open.defaults
@@ -33,16 +35,21 @@ def __init__(self, file_path: str, **options):
 
         # Remove uproot-specific options (should be done earlier)
         exclude_keys = set(default_options.keys())
-        opts = {k: v for k, v in options.items() if k not in exclude_keys}
+        storage_options = {k: v for k, v in options.items() if k not in exclude_keys}
 
-        self._fs, self._file_path = fsspec.core.url_to_fs(file_path, **opts)
+        protocol = fsspec.core.split_protocol(file_path)[0]
+        fs_has_async_impl = fsspec.get_filesystem_class(protocol=protocol).async_impl
 
-        if self._use_threads:
+        if not self._use_threads:
+            self._executor = uproot.source.futures.TrivialExecutor()
+        elif fs_has_async_impl:
+            self._executor = FSSpecLoopExecutor(fsspec.asyn.get_loop())
+        else:
             self._executor = concurrent.futures.ThreadPoolExecutor(
                 max_workers=self._num_workers
             )
-        else:
-            self._executor = uproot.source.futures.TrivialExecutor()
+
+        self._fs, self._file_path = fsspec.core.url_to_fs(file_path, **storage_options)
 
         # TODO: set mode to "read-only" in a way that works for all filesystems
         self._file = self._fs.open(self._file_path)
@@ -73,8 +80,8 @@ def __enter__(self):
 
     def __exit__(self, exception_type, exception_value, traceback):
         self._fh = None
-        self._executor.shutdown()
         self._file.__exit__(exception_type, exception_value, traceback)
+        self._executor.shutdown()
 
     def chunk(self, start: int, stop: int) -> uproot.source.chunk.Chunk:
         """
@@ -132,10 +139,11 @@ def chunks(
         self._num_requested_bytes += sum(stop - start for start, stop in ranges)
 
         chunks = []
+        # _cat_file is async while cat_file is not
+        use_async = isinstance(self._executor, FSSpecLoopExecutor)
+        cat_file = self._fs._cat_file if use_async else self._fs.cat_file
         for start, stop in ranges:
-            future = self._executor.submit(
-                self._fs.cat_file, self._file_path, start, stop
-            )
+            future = self._executor.submit(cat_file, self._file_path, start, stop)
             chunk = uproot.source.chunk.Chunk(self, start, stop, future)
             future.add_done_callback(uproot.source.chunk.notifier(chunk, notifications))
             chunks.append(chunk)
@@ -155,3 +163,16 @@ def closed(self) -> bool:
         otherwise.
         """
         return False
+
+
+class FSSpecLoopExecutor(uproot.source.futures.Executor):
+    def __init__(self, loop: asyncio.AbstractEventLoop):
+        self.loop = loop
+
+    def submit(self, coroutine, *args) -> asyncio.Future:
+        if not asyncio.iscoroutinefunction(coroutine):
+            raise TypeError("loop executor can only submit coroutines")
+        if not self.loop.is_running():
+            raise RuntimeError("cannot submit coroutine while loop is not running")
+        coroutine_object = coroutine(*args)
+        return asyncio.run_coroutine_threadsafe(coroutine_object, self.loop)

src/uproot/source/futures.py

@@ -28,6 +28,7 @@
 import sys
 import threading
 import time
+from abc import ABC, abstractmethod
 
 
 def delayed_raise(exception_class, exception_value, traceback):
@@ -37,6 +38,25 @@ def delayed_raise(exception_class, exception_value, traceback):
     raise exception_value.with_traceback(traceback)
 
 
+class Executor(ABC):
+    def __repr__(self):
+        return f"<{self.__class__.__name__} at 0x{id(self):012x}>"
+
+    @abstractmethod
+    def submit(self, task, *args):
+        """
+        Submit a task to be run in the background and return a Future object
+        representing that task.
+        """
+        raise NotImplementedError
+
+    def shutdown(self, wait: bool = True):
+        """
+        Stop the executor and free its resources.
+        """
+        return
+
+
 ##################### use-case 1: trivial Futures/Executor (satisfying formalities)
 
 
@@ -62,28 +82,20 @@ def result(self, timeout=None):
         return self._result
 
 
-class TrivialExecutor:
+class TrivialExecutor(Executor):
     """
     Formally satisfies the interface for a
     :doc:`uproot.source.futures.ThreadPoolExecutor`, but the
     :ref:`uproot.source.futures.TrivialExecutor.submit` method computes its
     ``task`` synchronously.
     """
 
-    def __repr__(self):
-        return f"<TrivialExecutor at 0x{id(self):012x}>"
-
     def submit(self, task, *args):
         """
         Immediately runs ``task(*args)``.
         """
         return TrivialFuture(task(*args))
 
-    def shutdown(self, wait: bool = True):
-        """
-        Does nothing, since this object does not have threads to stop.
-        """
-
 
 ##################### use-case 2: Python-like Futures/Executor for compute
 
@@ -176,7 +188,7 @@ def run(self):
             future._run()
 
 
-class ThreadPoolExecutor:
+class ThreadPoolExecutor(Executor):
     """
     Args:
         max_workers (None or int): The maximum number of workers to start.
@@ -207,9 +219,7 @@ def __init__(self, max_workers: int | None = None):
             worker.start()
 
     def __repr__(self):
-        return "<ThreadPoolExecutor ({} workers) at 0x{:012x}>".format(
-            len(self._workers), id(self)
-        )
+        return f"<{self.__class__.__name__} ({len(self._workers)} workers) at 0x{id(self):012x}>"
 
     @property
     def max_workers(self) -> int:
@@ -317,7 +327,7 @@ class ResourceWorker(Worker):
     executes.
     """
 
-    def __init__(self, work_queue, resource):
+    def __init__(self, work_queue: queue.Queue, resource):
         super().__init__(work_queue)
         self._resource = resource
 
@@ -368,11 +378,6 @@ def __init__(self, resources):
         for worker in self._workers:
             worker.start()
 
-    def __repr__(self):
-        return "<ResourceThreadPoolExecutor ({} workers) at 0x{:012x}>".format(
-            len(self._workers), id(self)
-        )
-
     def submit(self, future):
         """
         Pass the ``task`` onto the workers'
@@ -433,10 +438,7 @@ def __init__(self, resource):
         self._resource = resource
         self._closed = False
 
-    def __repr__(self):
-        return f"<ResourceTrivialExecutor at 0x{id(self):012x}>"
-
-    def submit(self, future):
+    def submit(self, future: ResourceFuture) -> ResourceFuture:
         """
         Pass the ``task`` as a
         :doc:`uproot.source.futures.ResourceFuture` so that it will be

tests/test_0692_fsspec.py

@@ -64,14 +64,22 @@ def test_open_fsspec_s3():
 @pytest.mark.network
 @pytest.mark.xrootd
 @pytest.mark.parametrize(
-    "handler",
-    [uproot.source.fsspec.FSSpecSource, uproot.source.xrootd.XRootDSource, None],
+    "handler, use_threads",
+    [
+        (uproot.source.fsspec.FSSpecSource, True),
+        (uproot.source.fsspec.FSSpecSource, False),
+        (uproot.source.xrootd.XRootDSource, True),
+        (uproot.source.xrootd.XRootDSource, False),
+        (None, True),
+        (None, False),
+    ],
 )
-def test_open_fsspec_xrootd(handler):
+def test_open_fsspec_xrootd(handler, use_threads):
     pytest.importorskip("XRootD")
     with uproot.open(
         "root://eospublic.cern.ch//eos/root-eos/cms_opendata_2012_nanoaod/Run2012B_DoubleMuParked.root",
         handler=handler,
+        use_threads=use_threads,
     ) as f:
         data = f["Events/run"].array(library="np", entry_stop=20)
         assert len(data) == 20