feat: asyncio LoopExecutor and async fsspec source

src/uproot/source/chunk.py

@@ -60,7 +60,9 @@ def chunk(self, start, stop) -> Chunk:
         :doc:`uproot.source.chunk.Chunk`.
         """
 
-    def chunks(self, ranges, notifications: queue.Queue) -> list[Chunk]:
+    def chunks(
+        self, ranges: list[(int, int)], notifications: queue.Queue
+    ) -> list[Chunk]:
         """
         Args:
             ranges (list of (int, int) 2-tuples): Intervals to fetch
@@ -164,7 +166,9 @@ def chunk(self, start, stop) -> Chunk:
         self._executor.submit(future)
         return chunk
 
-    def chunks(self, ranges, notifications: queue.Queue) -> list[Chunk]:
+    def chunks(
+        self, ranges: list[(int, int)], notifications: queue.Queue
+    ) -> list[Chunk]:
         self._num_requests += 1
         self._num_requested_chunks += len(ranges)
         self._num_requested_bytes += sum(stop - start for start, stop in ranges)

src/uproot/source/cursor.py

@@ -6,6 +6,7 @@
 the lowest level of interpretation (numbers, strings, raw arrays, etc.).
 """
 
+from __future__ import annotations
 
 import datetime
 import struct

src/uproot/source/file.py

@@ -156,7 +156,7 @@ def chunk(self, start, stop) -> uproot.source.chunk.Chunk:
             return self._fallback.chunk(start, stop)
 
     def chunks(
-        self, ranges, notifications: queue.Queue
+        self, ranges: list[(int, int)], notifications: queue.Queue
     ) -> list[uproot.source.chunk.Chunk]:
         if self._fallback is None:
             if self.closed:

src/uproot/source/fsspec.py

@@ -38,9 +38,22 @@ def __init__(self, file_path, **options):
         self._fs, self._file_path = fsspec.core.url_to_fs(file_path, **opts)
 
         if self._use_threads:
-            self._executor = concurrent.futures.ThreadPoolExecutor(
-                max_workers=self._num_workers
-            )
+            if self._fs.async_impl:
+                self._executor = uproot.source.futures.LoopExecutor()
+
+                # Bind the loop to the filesystem
+                async def make_fs():
+                    return fsspec.filesystem(
+                        protocol=self._fs.protocol, loop=self._executor.loop
+                    )
+
+                self._fs = self._executor.submit(make_fs).result()
+                assert self._fs.loop is self._executor.loop, "loop not bound"
+                assert self._fs.loop.is_running(), "loop not running"
+            else:
+                self._executor = concurrent.futures.ThreadPoolExecutor(
+                    max_workers=self._num_workers
+                )
         else:
             self._executor = uproot.source.futures.TrivialExecutor()
 
@@ -73,8 +86,11 @@ 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)
+        # TODO: proper cleanup. Shutting down the loop executor causes problems
+        # self._executor.shutdown()
+        if not isinstance(self._executor, uproot.source.futures.LoopExecutor):
+            self._executor.shutdown()
 
     def chunk(self, start, stop) -> uproot.source.chunk.Chunk:
         """
@@ -98,7 +114,7 @@ def chunk(self, start, stop) -> uproot.source.chunk.Chunk:
         return uproot.source.chunk.Chunk(self, start, stop, future)
 
     def chunks(
-        self, ranges, notifications: queue.Queue
+        self, ranges: list[(int, int)], notifications: queue.Queue
     ) -> list[uproot.source.chunk.Chunk]:
         """
         Args:
@@ -132,10 +148,13 @@ 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 = self._fs.async_impl and isinstance(
+            self._executor, uproot.source.futures.LoopExecutor
+        )
+        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)

src/uproot/source/futures.py

@@ -21,6 +21,9 @@
 These classes implement a *subset* of Python's Future and Executor interfaces.
 """
 
+from __future__ import annotations
+
+import asyncio
 import os
 import queue
 import sys
@@ -224,7 +227,7 @@ def num_workers(self) -> int:
         return len(self._workers)
 
     @property
-    def workers(self):
+    def workers(self) -> list[Worker]:
         """
         A list of workers (:doc:`uproot.source.futures.Worker`).
         """
@@ -315,7 +318,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
 
@@ -469,3 +472,45 @@ def __exit__(self, exception_type, exception_value, traceback):
         self.shutdown()
         self._resource.__exit__(exception_type, exception_value, traceback)
         self._closed = True
+
+
+class LoopExecutor:
+    def __repr__(self):
+        return f"<LoopExecutor at 0x{id(self):012x}>"
+
+    def __init__(self):
+        self._loop = asyncio.new_event_loop()
+        # TODO: remove daemon=True (or not?)
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self.start()
+
+    def start(self):
+        self._thread.start()
+        return self
+
+    def shutdown(self):
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        self._thread.join()
+
+    def _run(self):
+        asyncio.set_event_loop(self._loop)
+        try:
+            self._loop.run_forever()
+        finally:
+            self._loop.run_until_complete(self._loop.shutdown_asyncgens())
+            self._loop.close()
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.shutdown()
+
+    @property
+    def loop(self) -> asyncio.AbstractEventLoop:
+        return self._loop
+
+    def submit(self, coroutine, *args) -> asyncio.Future:
+        coroutine_object = coroutine(*args)
+        return asyncio.run_coroutine_threadsafe(coroutine_object, self._loop)

src/uproot/source/http.py

@@ -260,7 +260,7 @@ def task(resource):
         return uproot.source.futures.ResourceFuture(task)
 
     @staticmethod
-    def multifuture(source, ranges, futures, results):
+    def multifuture(source, ranges: list[(int, int)], futures, results):
         """
         Args:
             source (:doc:`uproot.source.http.HTTPSource`): The data source.
@@ -351,7 +351,7 @@ def task(resource):
     )
     _content_range = re.compile(b"Content-Range: bytes ([0-9]+-[0-9]+)", re.I)
 
-    def is_multipart_supported(self, ranges, response):
+    def is_multipart_supported(self, ranges: list[(int, int)], response) -> bool:
         """
         Helper function for :ref:`uproot.source.http.HTTPResource.multifuture`
         to check for multipart GET support.
@@ -368,7 +368,7 @@ def is_multipart_supported(self, ranges, response):
         else:
             return True
 
-    def handle_no_multipart(self, source, ranges, futures, results):
+    def handle_no_multipart(self, source, ranges: list[(int, int)], futures, results):
         """
         Helper function for :ref:`uproot.source.http.HTTPResource.multifuture`
         to handle a lack of multipart GET support.
@@ -605,7 +605,7 @@ def chunk(self, start, stop) -> uproot.source.chunk.Chunk:
         return chunk
 
     def chunks(
-        self, ranges, notifications: queue.Queue
+        self, ranges: list[(int, int)], notifications: queue.Queue
     ) -> list[uproot.source.chunk.Chunk]:
         if self._fallback is None:
             self._num_requests += 1

src/uproot/source/object.py

@@ -8,6 +8,7 @@
 has exactly one worker (we can't assume that the object is thread-safe).
 """
 
+from __future__ import annotations
 
 import uproot
 import uproot.source.chunk

src/uproot/source/s3.py

@@ -4,6 +4,8 @@
 This module defines a physical layer for remote files, accessed via S3.
 """
 
+from __future__ import annotations
+
 import os
 from urllib.parse import parse_qsl, urlparse
 

src/uproot/source/xrootd.py

@@ -325,7 +325,7 @@ def chunk(self, start, stop) -> uproot.source.chunk.Chunk:
         return uproot.source.chunk.Chunk(self, start, stop, future)
 
     def chunks(
-        self, ranges, notifications: queue.Queue
+        self, ranges: list[(int, int)], notifications: queue.Queue
     ) -> list[uproot.source.chunk.Chunk]:
         self._num_requests += 1
         self._num_requested_chunks += len(ranges)