Repository: hikari-py/hikari
Pull request #1272 (2022-08-26 19:22:25)

Add ability to save resources
### Summary
This PR adds the ability to save attachments (and more generally, resources) to disk.

* Adds the `Resource.save()` method
* Adds tests for `Resource.save()` and `Resource.read()`

### Checklist
- [x] I have run `nox` and all the pipelines have passed.
- [x] I have made unittests according to the code I have added/modified/deleted.


### Related issues
None


davfsa: > Honestly I think I'm just gonna have to leave the tests up to somebody else to take care of cos I can't figure this out. I've emulated what you've done, and I'm getting all sorts of errors and have absolutely no idea what I'm doing wrong.

@parafoxia mind pushing what you have (can be with broken tests too, or just skip them entirely).
parafoxia: @FasterSpeeding, just bumping this as it's ready for review.
FasterSpeeding: Test cases for the specialisations?
davfsa: > Test cases for the specialisations?

Working on those atm :)
davfsa: @FasterSpeeding tests done. I also found some bugs in the code, so some parts are update too :)

Patch:
diff --git a/changes/1272.feature.md b/changes/1272.feature.md
new file mode 100644
index 00000000..8d755dce
--- /dev/null
+++ b/changes/1272.feature.md
@@ -0,0 +1,1 @@
+Add `save()` method to `Resource`.
diff --git a/hikari/files.py b/hikari/files.py
index c313fbd3..2e84d062 100644
--- a/hikari/files.py
+++ b/hikari/files.py
@@ -52,6 +52,7 @@
 import mimetypes
 import os
 import pathlib
+import shutil
 import typing
 import urllib.parse
 import urllib.request
@@ -408,6 +409,22 @@ async def __aexit__(
         pass
 
 
+def _to_write_path(path: Pathish, default_filename: str, force: bool) -> pathlib.Path:
+    path = ensure_path(path)
+    if path.is_dir():
+        path = path.joinpath(default_filename)
+
+    if not force and path.exists():
+        raise FileExistsError(f"file {path!r} already exists; use `force=True` to overwrite")
+
+    return path.expanduser()
+
+
+def _open_write_path(path: Pathish, default_filename: str, force: bool) -> typing.BinaryIO:
+    path = _to_write_path(path, default_filename, force)
+    return path.open("wb")
+
+
 class Resource(typing.Generic[ReaderImplT], abc.ABC):
     """Base for any uploadable or downloadable representation of information.
 
@@ -470,6 +487,40 @@ async def read(
         async with self.stream(executor=executor) as reader:
             return await reader.read()
 
+    async def save(
+        self,
+        path: Pathish,
+        *,
+        executor: typing.Optional[concurrent.futures.Executor] = None,
+        force: bool = False,
+    ) -> None:
+        """Save this resource to disk.
+
+        This writes the resource file in chunks, and so does not load
+        the entire resource into memory.
+
+        Parameters
+        ----------
+        path : Pathish
+            The path to save this resource to. If this is a string, the
+            path will be relative to the current working directory.
+        executor : typing.Optional[concurrent.futures.Executor]
+            The executor to run in for blocking operations.
+            If `None`, then the default executor is used for
+            the current event loop.
+        force : bool
+            Whether to overwrite an existing file. Defaults to `False`.
+        """
+        loop = asyncio.get_running_loop()
+        file = await loop.run_in_executor(executor, _open_write_path, path, self.filename, force)
+
+        try:
+            async with self.stream(executor=executor) as reader:
+                async for chunk in reader:
+                    await loop.run_in_executor(executor, file.write, chunk)
+        finally:
+            await loop.run_in_executor(executor, file.close)
+
     @abc.abstractmethod
     def stream(
         self,
@@ -784,7 +835,7 @@ async def __aiter__(self) -> typing.AsyncGenerator[typing.Any, bytes]:
                 break
 
 
-def _open_path(path: pathlib.Path) -> typing.BinaryIO:
+def _open_read_path(path: pathlib.Path) -> typing.BinaryIO:
     return path.expanduser().open("rb")
 
 
@@ -801,7 +852,7 @@ async def __aenter__(self) -> ThreadedFileReader:
             raise RuntimeError("File is already open")
 
         loop = asyncio.get_running_loop()
-        file = await loop.run_in_executor(self.executor, _open_path, self.path)
+        file = await loop.run_in_executor(self.executor, _open_read_path, self.path)
         self.file = file
         return ThreadedFileReader(self.filename, None, self.executor, file)
 
@@ -819,6 +870,11 @@ async def __aexit__(
         self.file = None
 
 
+def _copy_to_path(current_path: pathlib.Path, copy_to_path: Pathish, default_filename: str, force: bool) -> None:
+    copy_to_path = _to_write_path(copy_to_path, default_filename, force)
+    shutil.copy2(current_path, copy_to_path)
+
+
 class File(Resource[ThreadedFileReader]):
     """A resource that exists on the local machine's storage to be uploaded.
 
@@ -908,6 +964,18 @@ def stream(
 
         raise TypeError("The executor must be a ThreadPoolExecutor or None")
 
+    async def save(
+        self,
+        path: Pathish,
+        *,
+        executor: typing.Optional[concurrent.futures.Executor] = None,
+        force: bool = False,
+    ) -> None:
+        # An optimization can be done here to avoid a lot of thread calls and streaming
+        # by just copying the file
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(executor, _copy_to_path, self.path, path, self.filename, force)
+
 
 ########################################################################
 # RAW BYTE, ASYNC ITERATOR, ASYNC ITERABLE, ITERATOR, ITERABLE READERS #
@@ -986,6 +1054,13 @@ def _assert_bytes(data: typing.Any) -> bytes:
         return data
 
 
+def _write_bytes(
+    path: Pathish, default_filename: str, force: bool, data: typing.Union[bytearray, bytes, memoryview]
+) -> None:
+    path = _to_write_path(path, default_filename, force)
+    path.write_bytes(data)
+
+
 class Bytes(Resource[IteratorReader]):
     """Representation of in-memory data to upload.
 
@@ -1029,10 +1104,7 @@ def __init__(
         self.data = data
 
         if mimetype is None:
-            mimetype = guess_mimetype_from_filename(filename)
-
-        if mimetype is None:
-            mimetype = "text/plain;charset=UTF-8"
+            mimetype = guess_mimetype_from_filename(filename) or "text/plain;charset=UTF-8"
 
         self._filename = filename
         self.mimetype = mimetype
@@ -1072,6 +1144,22 @@ def stream(
         """
         return _NoOpAsyncReaderContextManagerImpl(IteratorReader(self.filename, self.mimetype, self.data))
 
+    async def save(
+        self,
+        path: Pathish,
+        *,
+        executor: typing.Optional[concurrent.futures.Executor] = None,
+        force: bool = False,
+    ) -> None:
+        if not isinstance(self.data, (bytes, bytearray, memoryview)):
+            await super().save(path, executor=executor, force=force)
+            return
+
+        # An optimization can be done here to avoid a lot of thread calls and streaming
+        # by just writing the whole data at once
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(executor, _write_bytes, path, self.filename, force, self.data)
+
     @staticmethod
     def from_data_uri(data_uri: str, filename: typing.Optional[str] = None) -> Bytes:
         """Parse a given data URI.


Test patch:
diff --git a/tests/hikari/test_files.py b/tests/hikari/test_files.py
index 91ad4e0d..6661aaa3 100644
--- a/tests/hikari/test_files.py
+++ b/tests/hikari/test_files.py
@@ -22,6 +22,7 @@
 
 import asyncio
 import pathlib
+import shutil
 
 import mock
 import pytest
@@ -59,11 +60,11 @@ def test___exit__(self, reader):
             pytest.fail(exc)
 
 
-def test__open_path():
+def test__open_read_path():
     expanded_path = mock.Mock()
     path = mock.Mock(expanduser=mock.Mock(return_value=expanded_path))
 
-    assert files._open_path(path) is expanded_path.open.return_value
+    assert files._open_read_path(path) is expanded_path.open.return_value
 
     expanded_path.open.assert_called_once_with("rb")
 
@@ -100,8 +101,175 @@ async def test_context_manager(self):
                 assert reader.filename == "meow.txt"
                 assert reader._executor is executor
 
-                loop.run_in_executor.assert_called_once_with(executor, files._open_path, path)
+                loop.run_in_executor.assert_called_once_with(executor, files._open_read_path, path)
                 loop.run_in_executor.reset_mock()
 
         loop.run_in_executor.assert_called_once_with(executor, mock_file.close)
         assert context_manager.file is None
+
+
+class TestToWritePath:
+    def test_when_dir(self):
+        mock_path = mock.Mock(is_dir=mock.Mock(return_value=True))
+        joinpath = mock_path.joinpath.return_value = mock.Mock(exists=mock.Mock(return_value=False))
+
+        with mock.patch.object(files, "ensure_path", return_value=mock_path) as ensure_path:
+            assert files._to_write_path("test_path", "some_filename.png", False) is joinpath.expanduser.return_value
+
+        mock_path.joinpath.assert_called_once_with("some_filename.png")
+        ensure_path.assert_called_once_with("test_path")
+
+    def test_when_exists_but_not_force(self):
+        mock_path = mock.Mock(is_dir=mock.Mock(return_value=False), exists=mock.Mock(return_value=True))
+
+        with mock.patch.object(files, "ensure_path", return_value=mock_path) as ensure_path:
+            with pytest.raises(FileExistsError):
+                files._to_write_path("test_path", "some_filename.png", False)
+
+        ensure_path.assert_called_once_with("test_path")
+
+    def test_when_exists_but_force(self):
+        mock_path = mock.Mock(is_dir=mock.Mock(return_value=False), exists=mock.Mock(return_value=True))
+
+        with mock.patch.object(files, "ensure_path", return_value=mock_path) as ensure_path:
+            assert files._to_write_path("test_path", "some_filename.png", True) is mock_path.expanduser.return_value
+
+        ensure_path.assert_called_once_with("test_path")
+
+
+def test_open_write_path():
+    with mock.patch.object(files, "_to_write_path") as to_write_path:
+        assert (
+            files._open_write_path("path", "some_filename.png", False) is to_write_path.return_value.open.return_value
+        )
+
+    to_write_path.assert_called_once_with("path", "some_filename.png", False)
+    to_write_path.return_value.open.assert_called_once_with("wb")
+
+
+class TestResource:
+    @pytest.fixture()
+    def resource(self):
+        class MockReader:
+            data = iter(("never", "gonna", "give", "you", "up"))
+
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, *args, **kwargs):
+                return
+
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                try:
+                    return next(self.data)
+                except StopIteration:
+                    raise StopAsyncIteration from None
+
+        class ResourceImpl(files.Resource):
+            stream = mock.Mock(return_value=MockReader())
+            url = "https://myspace.com/rickastley/lyrics.txt"
+            filename = "lyrics.txt"
+
+        return ResourceImpl()
+
+    @pytest.mark.asyncio()
+    async def test_save(self, resource):
+        executor = object()
+        file_open = mock.Mock()
+        file_open.write = mock.Mock()
+        loop = mock.Mock(run_in_executor=mock.AsyncMock(side_effect=[file_open, None, None, None, None, None, None]))
+
+        with mock.patch.object(asyncio, "get_running_loop", return_value=loop):
+            await resource.save("rickroll/lyrics.txt", executor=executor, force=True)
+
+        assert loop.run_in_executor.call_count == 7
+        loop.run_in_executor.assert_has_calls(
+            [
+                mock.call(executor, files._open_write_path, "rickroll/lyrics.txt", "lyrics.txt", True),
+                mock.call(executor, file_open.write, "never"),
+                mock.call(executor, file_open.write, "gonna"),
+                mock.call(executor, file_open.write, "give"),
+                mock.call(executor, file_open.write, "you"),
+                mock.call(executor, file_open.write, "up"),
+                mock.call(executor, file_open.close),
+            ]
+        )
+
+
+def test_copy_to_path():
+    with mock.patch.object(files, "_to_write_path") as to_write_path:
+        with mock.patch.object(shutil, "copy2") as copy2:
+            files._copy_to_path("original_path", "path", "some_filename.png", False)
+
+    to_write_path.assert_called_once_with("path", "some_filename.png", False)
+    copy2.assert_called_once_with("original_path", to_write_path.return_value)
+
+
+class TestFile:
+    @pytest.fixture()
+    def file_obj(self):
+        return files.File("one/path/something.txt")
+
+    @pytest.mark.asyncio()
+    async def test_save(self, file_obj):
+        mock_executor = object()
+        loop = mock.Mock(run_in_executor=mock.AsyncMock())
+
+        with mock.patch.object(asyncio, "get_running_loop", return_value=loop):
+            with mock.patch.object(files.Resource, "save") as super_save:
+                await file_obj.save("some_path/", executor=mock_executor, force=True)
+
+        super_save.assert_not_called()
+        loop.run_in_executor.assert_awaited_once_with(
+            mock_executor,
+            files._copy_to_path,
+            pathlib.Path("one/path/something.txt"),
+            "some_path/",
+            "something.txt",
+            True,
+        )
+
+
+def test_write_bytes():
+    with mock.patch.object(files, "_to_write_path") as to_write_path:
+        files._write_bytes("path", "some_filename.png", False, b"some bytes")
+
+    to_write_path.assert_called_once_with("path", "some_filename.png", False)
+    to_write_path.return_value.write_bytes.assert_called_once_with(b"some bytes")
+
+
+class TestBytes:
+    @pytest.fixture()
+    def bytes_obj(self):
+        return files.Bytes(b"some data", "something.txt")
+
+    @pytest.mark.parametrize("data_type", [bytes, bytearray, memoryview])
+    @pytest.mark.asyncio()
+    async def test_save(self, bytes_obj, data_type):
+        bytes_obj.data = mock.Mock(data_type)
+        mock_executor = object()
+        loop = mock.Mock(run_in_executor=mock.AsyncMock())
+
+        with mock.patch.object(asyncio, "get_running_loop", return_value=loop):
+            with mock.patch.object(files.Resource, "save") as super_save:
+                await bytes_obj.save("some_path/", executor=mock_executor, force=True)
+
+        super_save.assert_not_called()
+        loop.run_in_executor.assert_awaited_once_with(
+            mock_executor, files._write_bytes, "some_path/", "something.txt", True, bytes_obj.data
+        )
+
+    @pytest.mark.asyncio()
+    async def test_save_when_data_is_not_bytes(self, bytes_obj):
+        bytes_obj.data = object()
+        mock_executor = object()
+
+        with mock.patch.object(asyncio, "get_running_loop") as get_running_loop:
+            with mock.patch.object(files.Resource, "save") as super_save:
+                await bytes_obj.save("some_path/", executor=mock_executor, force=True)
+
+        super_save.assert_called_once_with("some_path/", executor=mock_executor, force=True)
+        get_running_loop.return_value.run_in_executor.assert_not_called()