feat: Add Async File Support (#34)

Author: Xuanwo

opendalfs/file.py

@@ -1,6 +1,8 @@
 import logging
 
+from fsspec.asyn import AbstractAsyncStreamedFile
 from fsspec.spec import AbstractBufferedFile
+from opendal import AsyncFile as OpendalAsyncFile
 from opendal import File as OpendalFile
 
 logger = logging.getLogger("opendalfs")
@@ -148,3 +150,133 @@ def close(self):
                     self._opendal_writer.close()
                 finally:
                     self._opendal_writer = None
+
+
+class OpendalAsyncBufferedFile(AbstractAsyncStreamedFile):
+    """Async buffered file implementation for OpenDAL."""
+
+    _opendal_writer: OpendalAsyncFile | None
+    _append_via_write: bool
+    _initiated: bool
+    _exclusive_create: bool
+
+    def __init__(
+        self,
+        fs,
+        path,
+        mode="rb",
+        block_size="default",
+        autocommit=True,
+        cache_type="readahead",
+        cache_options=None,
+        size=None,
+        **kwargs,
+    ):
+        self._exclusive_create = mode == "xb"
+        normalized_mode = "wb" if self._exclusive_create else mode
+        super().__init__(
+            fs,
+            path,
+            mode=normalized_mode,
+            block_size=block_size,
+            autocommit=autocommit,
+            cache_type=cache_type,
+            cache_options=cache_options,
+            size=size,
+            **kwargs,
+        )
+
+        self._opendal_writer = None
+        self._append_via_write = False
+        self._initiated = False
+
+    async def _fetch_range(self, start: int, end: int):
+        if start >= end:
+            return b""
+
+        reader = await self.fs.async_fs.open(self.path, "rb")
+        try:
+            await reader.seek(start)
+            return await reader.read(end - start)
+        finally:
+            await reader.close()
+
+    async def _upload_chunk(self, final: bool = False):
+        if not self._initiated:
+            raise RuntimeError("Upload has not been initiated")
+
+        self.buffer.seek(0)
+        chunk = self.buffer.read()
+
+        if not chunk:
+            if not final:
+                return False
+            if self.mode == "ab" and self._append_via_write:
+                if not await self.fs.async_fs.exists(self.path):
+                    await self.fs.async_fs.write(self.path, b"")
+                return None
+            await self._commit_upload()
+            return None
+
+        if self.mode == "ab" and self._append_via_write:
+            await self.fs.async_fs.write(self.path, chunk, append=True)
+            return None
+
+        if self._opendal_writer is None:
+            self._opendal_writer = await self.fs.async_fs.open(self.path, "wb")
+
+        await self._opendal_writer.write(chunk)
+
+        if final:
+            await self._commit_upload()
+        return None
+
+    async def _initiate_upload(self) -> None:
+        if self._initiated:
+            return
+
+        if self._exclusive_create and await self.fs.async_fs.exists(self.path):
+            raise FileExistsError(self.path)
+
+        if self.mode == "ab":
+            cap = self.fs.async_fs.capability()
+            if getattr(cap, "write_can_append", False):
+                self._append_via_write = True
+                self.offset = self.loc
+            else:
+                try:
+                    existing = await self.fs.async_fs.read(self.path)
+                except FileNotFoundError:
+                    existing = b""
+                if existing:
+                    self._opendal_writer = await self.fs.async_fs.open(self.path, "wb")
+                    await self._opendal_writer.write(existing)
+                    self.offset = len(existing)
+
+        self._initiated = True
+
+    async def _commit_upload(self) -> None:
+        if self.mode == "ab" and self._append_via_write:
+            return
+
+        if self._opendal_writer is None:
+            await self.fs.async_fs.write(self.path, b"")
+            return
+
+        try:
+            await self._opendal_writer.close()
+        finally:
+            self._opendal_writer = None
+
+    async def close(self):
+        if self.closed:
+            return
+
+        try:
+            await super().close()
+        finally:
+            if self._opendal_writer is not None:
+                try:
+                    await self._opendal_writer.close()
+                finally:
+                    self._opendal_writer = None

opendalfs/fs.py

@@ -2,7 +2,7 @@
 from fsspec.asyn import AsyncFileSystem
 import logging
 from opendal import AsyncOperator, Operator
-from .file import OpendalBufferedFile
+from .file import OpendalAsyncBufferedFile, OpendalBufferedFile
 
 logger = logging.getLogger("opendalfs")
 
@@ -147,6 +147,26 @@ def _open(
             **kwargs,
         )
 
+    async def open_async(self, path, mode="rb", **kwargs):
+        if "b" not in mode or kwargs.get("compression"):
+            raise ValueError
+
+        size = None
+        if mode == "rb":
+            info = await self.async_fs.stat(path)
+            size = info.content_length
+
+        file = OpendalAsyncBufferedFile(self, path, mode, size=size, **kwargs)
+
+        if mode == "ab":
+            try:
+                info = await self.async_fs.stat(path)
+                file.loc = info.content_length
+            except FileNotFoundError:
+                file.loc = 0
+
+        return file
+
     async def _modified(self, path: str):
         """Get modified time (async version)"""
         info = await self.async_fs.stat(path)

tests/core/test_file.py

@@ -31,3 +31,62 @@ def test_open_exclusive_create(any_fs):
     with pytest.raises(FileExistsError):
         with any_fs.open("exists.txt", "xb") as f:
             f.write(b"y")
+
+
+@pytest.mark.asyncio
+async def test_open_async_read_seek():
+    from opendalfs import OpendalFileSystem
+
+    fs = OpendalFileSystem(scheme="memory", asynchronous=True, skip_instance_cache=True)
+    data = b"0123456789"
+    await fs._pipe_file("readseek.txt", data)
+
+    async with await fs.open_async("readseek.txt", "rb") as f:
+        assert await f.read(3) == b"012"
+        assert f.tell() == 3
+
+        f.seek(5)
+        assert await f.read(2) == b"56"
+
+        f.seek(-3, 2)
+        assert await f.read() == b"789"
+
+
+@pytest.mark.asyncio
+async def test_open_async_write_chunked():
+    from opendalfs import OpendalFileSystem
+
+    fs = OpendalFileSystem(scheme="memory", asynchronous=True, skip_instance_cache=True)
+
+    async with await fs.open_async("chunked.txt", "wb", block_size=3) as f:
+        await f.write(b"abc")
+        await f.write(b"def")
+        await f.write(b"gh")
+
+    assert await fs._cat_file("chunked.txt") == b"abcdefgh"
+
+
+@pytest.mark.asyncio
+async def test_open_async_exclusive_create():
+    from opendalfs import OpendalFileSystem
+
+    fs = OpendalFileSystem(scheme="memory", asynchronous=True, skip_instance_cache=True)
+    await fs._pipe_file("exists.txt", b"x")
+
+    with pytest.raises(FileExistsError):
+        async with await fs.open_async("exists.txt", "xb") as f:
+            await f.write(b"y")
+
+
+@pytest.mark.asyncio
+async def test_open_async_append_emulated():
+    from opendalfs import OpendalFileSystem
+
+    fs = OpendalFileSystem(scheme="memory", asynchronous=True, skip_instance_cache=True)
+    await fs._pipe_file("append.txt", b"hello")
+
+    async with await fs.open_async("append.txt", "ab", block_size=2) as f:
+        assert f.tell() == 5
+        await f.write(b"world")
+
+    assert await fs._cat_file("append.txt") == b"helloworld"