Merge pull request #37 from fsspec/enable_simplecache

implement _get_file()

Author: lkluft

ipfsspec/async_ipfs.py

@@ -4,6 +4,7 @@
 import weakref
 from functools import lru_cache
 from pathlib import Path
+from contextlib import asynccontextmanager
 import warnings
 
 import asyncio
@@ -12,6 +13,8 @@
 
 from fsspec.asyn import AsyncFileSystem, sync, sync_wrapper
 from fsspec.exceptions import FSTimeoutError
+from fsspec.callbacks import DEFAULT_CALLBACK
+from fsspec.utils import isfilelike
 
 from multiformats import CID, multicodec
 from . import unixfsv1
@@ -118,6 +121,17 @@ async def cat(self, path, session):
             self._raise_not_found_for_status(res, path)
             return await res.read()
 
+    @asynccontextmanager
+    async def iter_chunked(self, path, session, chunk_size):
+        res = await self.get(path, session)
+        async with res:
+            self._raise_not_found_for_status(res, path)
+            try:
+                size = int(res.headers["content-length"])
+            except (ValueError, KeyError):
+                size = None
+            yield size, res.content.iter_chunked(chunk_size)
+
     async def ls(self, path, session, detail=False):
         res = await self.get(path, session, headers={"Accept": "application/vnd.ipld.raw"}, params={"format": "raw"})
         self._raise_not_found_for_status(res, path)
@@ -293,6 +307,28 @@ async def _cat_file(self, path, start=None, end=None, **kwargs):
         session = await self.set_session()
         return (await self.gateway.cat(path, session))[start:end]
 
+    async def _get_file(
+        self, rpath, lpath, chunk_size=5 * 2**20, callback=DEFAULT_CALLBACK, **kwargs
+    ):
+        logger.debug(rpath)
+        rpath = self._strip_protocol(rpath)
+        session = await self.set_session()
+
+        if isfilelike(lpath):
+            outfile = lpath
+        else:
+            outfile = open(lpath, "wb")  # noqa: ASYNC101, ASYNC230
+
+        try:
+            async with self.gateway.iter_chunked(rpath, session, chunk_size) as (size, chunks):
+                callback.set_size(size)
+                async for chunk in chunks:
+                    outfile.write(chunk)
+                    callback.relative_update(len(chunk))
+        finally:
+            if not isfilelike(lpath):
+                outfile.close()
+
     async def _info(self, path, **kwargs):
         path = self._strip_protocol(path)
         session = await self.set_session()

test/test_async.py

@@ -80,6 +80,12 @@ async def test_cat_file(fs):
     assert res == REF_CONTENT[3:7]
 
 
+@pytest.mark.asyncio
+async def test_get_file(fs, tmp_path):
+    await fs._get_file(TEST_ROOT + "/default", tmp_path / "default")
+    assert open(tmp_path / "default", "rb").read() == REF_CONTENT
+
+
 @pytest.mark.asyncio
 async def test_exists(fs):
     res = await fs._exists(TEST_ROOT + "/default")