From 49b7d0816ef37a342abde789e1a82d893db10333 Mon Sep 17 00:00:00 2001 From: skshetry <18718008+skshetry@users.noreply.github.com> Date: Sat, 13 Jan 2024 19:27:53 +0545 Subject: [PATCH] asynlocal: make open_async compatible with fsspec (#25) * asynlocal: make open_async compatible with fsspec Closes #24. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- src/morefs/asyn_local.py | 9 +++------ src/morefs/overlay.py | 2 +- tests/test_asyn_local.py | 15 ++++++++++----- 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src/morefs/asyn_local.py b/src/morefs/asyn_local.py index f547419..abc4b2c 100644 --- a/src/morefs/asyn_local.py +++ b/src/morefs/asyn_local.py @@ -1,7 +1,6 @@ import shutil import sys from asyncio import get_running_loop, iscoroutinefunction -from contextlib import asynccontextmanager from functools import partial, wraps from typing import Awaitable, Callable, TypeVar @@ -83,18 +82,16 @@ async def _get_file(self, src, dst, **kwargs): # pylint: disable=arguments-rena src = self._strip_protocol(src) return await self._get_file_async(src, dst) - async with self.open_async(src, "rb") as fsrc: + fsrc = await self.open_async(src, "rb") + async with fsrc: while True: buf = await fsrc.read(length=shutil.COPY_BUFSIZE) if not buf: break await dst.write(buf) - @asynccontextmanager async def open_async(self, path, mode="rb", **kwargs): path = self._strip_protocol(path) if self.auto_mkdir and "w" in mode: await self._makedirs(self._parent(path), exist_ok=True) - - async with aiofile.async_open(path, mode, **kwargs) as f: - yield f + return await aiofile.async_open(path, mode, **kwargs) diff --git a/src/morefs/overlay.py b/src/morefs/overlay.py index 7569a29..1ddda49 100644 --- a/src/morefs/overlay.py +++ b/src/morefs/overlay.py @@ -6,7 +6,7 @@ import fsspec -class OverlayFileSystem(fsspec.AbstractFileSystem): +class OverlayFileSystem(fsspec.AbstractFileSystem): # pylint: disable=abstract-method cachable = False def __init__(self, *fses: fsspec.AbstractFileSystem, **kwargs): diff --git a/tests/test_asyn_local.py b/tests/test_asyn_local.py index 3c220f7..f6fc0c8 100644 --- a/tests/test_asyn_local.py +++ b/tests/test_asyn_local.py @@ -78,14 +78,17 @@ def test_sync_methods(tmp_path, localfs, fs): @pytest.mark.asyncio async def test_open_async(tmp_path, fs): - async with fs.open_async(tmp_path / "file", mode="wb") as f: + f = await fs.open_async(tmp_path / "file", mode="wb") + async with f: pass assert await fs._exists(tmp_path / "file") - async with fs.open_async(tmp_path / "file", mode="wb") as f: + f = await fs.open_async(tmp_path / "file", mode="wb") + async with f: assert await f.write(b"contents") - async with fs.open_async(tmp_path / "file") as f: + f = await fs.open_async(tmp_path / "file") + async with f: assert await f.read() == b"contents" @@ -96,7 +99,8 @@ async def test_get_file(tmp_path, fs): assert await fs._isfile(tmp_path / "bar") - async with fs.open_async(tmp_path / "file1", mode="wb") as f: + f = await fs.open_async(tmp_path / "file1", mode="wb") + async with f: await fs._get_file(tmp_path / "foo", f) assert await fs._cat_file(tmp_path / "file1") == b"foo" @@ -112,7 +116,8 @@ async def test_get_file(tmp_path, fs): @pytest.mark.asyncio async def test_auto_mkdir_on_open_async(tmp_path): fs = AsyncLocalFileSystem(auto_mkdir=True) - async with fs.open_async(tmp_path / "dir" / "file", mode="wb") as f: + f = await fs.open_async(tmp_path / "dir" / "file", mode="wb") + async with f: await f.write(b"contents") assert await fs._isdir(tmp_path / "dir")