Add semaphore as argument to AsyncFileSystemWrapper

Author: dgegen

fsspec/implementations/asyn_wrapper.py

@@ -57,6 +57,8 @@ def __init__(
         asynchronous=None,
         target_protocol=None,
         target_options=None,
+        semaphore=None,
+        max_concurrent_tasks=None,
         **kwargs,
     ):
         if asynchronous is None:
@@ -67,7 +69,7 @@ def __init__(
         else:
             self.sync_fs = fsspec.filesystem(target_protocol, **target_options)
         self.protocol = self.sync_fs.protocol
-        self.semaphore = asyncio.Semaphore(1) if not asynchronous else None
+        self.semaphore = semaphore
         self._wrap_all_sync_methods()
 
     @property

fsspec/implementations/tests/test_asyn_wrapper.py

@@ -199,7 +199,9 @@ def test_open(tmpdir):
 
 @pytest.mark.asyncio
 async def test_semaphore_synchronous():
-    fs = AsyncFileSystemWrapper(LockedFileSystem(), asynchronous=False)
+    fs = AsyncFileSystemWrapper(
+            LockedFileSystem(), asynchronous=False, semaphore=asyncio.Semaphore(1)
+    )
 
     paths = [f"path_{i}" for i in range(1, 3)]
     results = await asyncio.gather(*(fs._cat_file(path) for path in paths))
@@ -209,7 +211,7 @@ async def test_semaphore_synchronous():
 
 @pytest.mark.asyncio
 async def test_deadlock_when_asynchronous():
-    fs = AsyncFileSystemWrapper(LockedFileSystem(), asynchronous=True)
+    fs = AsyncFileSystemWrapper(LockedFileSystem(), asynchronous=False, semaphore=asyncio.Semaphore(3))
     paths = [f"path_{i}" for i in range(1, 3)]
 
     with pytest.raises(RuntimeError, match="Concurrent requests!"):