block/upload: add download_file() (backport of jamalex#209)

Author: arturtamborski

notion/block/upload.py

@@ -1,10 +1,11 @@
 import os
 from mimetypes import guess_type
-from urllib.parse import urlencode, urlparse
+from urllib.parse import urlencode, urlparse, quote
 
 from notion.block.embed import EmbedBlock
 from notion.maps import field_map, property_map
-from notion.utils import human_size
+from notion.settings import SIGNED_URL_PREFIX
+from notion.utils import human_size, from_list
 
 
 class UploadBlock(EmbedBlock):
@@ -61,6 +62,45 @@ def upload_file(self, path: str):
         self.display_source = query_url
         self.file_id = urlparse(url).path.split("/")[2]
 
+    def download_file(self, path: str):
+        """
+        Download a file.
+
+
+        Arguments
+        ---------
+        path : str
+            Path for saving file.
+
+
+        Raises
+        ------
+        HTTPError
+            On API error.
+        """
+
+        record_data = self._get_record_data()
+        source = record_data["properties"]["source"]
+        s3_url = from_list(source)
+        file_name = s3_url.split("/")[-1]
+
+        params = {
+            "cache": "v2",
+            "name": file_name,
+            "id": self._id,
+            "table": self._table,
+            "userId": self._client.current_user.id,
+            "download": True,
+        }
+
+        url = SIGNED_URL_PREFIX + quote(s3_url, safe="")
+        resp = self._client.session.get(url, params=params, stream=True)
+        resp.raise_for_status()
+
+        with open(path, "wb") as f:
+            for chunk in resp.iter_content(chunk_size=4096):
+                f.write(chunk)
+
 
 class FileBlock(UploadBlock):
     """

notion/utils.py

@@ -33,6 +33,28 @@ def to_list(value) -> list:
     return value if isinstance(value, list) else [value]
 
 
+def from_list(value) -> Any:
+    """
+    Unwrap value from nested list.
+
+
+    Arguments
+    ---------
+    value : List
+        Nested list with target value.
+
+
+    Returns
+    -------
+    Any
+        Value from nested list.
+    """
+    if "__iter__" in dir(value) and not isinstance(value, str):
+        return from_list(next(iter(value), None))
+
+    return value
+
+
 def now() -> int:
     """
     Get UNIX-style time since epoch in seconds.

smoke_tests/block/test_upload.py

@@ -1,3 +1,5 @@
+import os
+
 from notion.block.upload import PdfBlock, ImageBlock, VideoBlock, FileBlock
 from smoke_tests.conftest import assert_block_is_okay, assert_block_attributes
 
@@ -19,19 +21,30 @@ def test_file_block(notion):
     assert "secure.notion-static.com" in block.source
     assert len(block.file_id) == 36
 
+    title_bak = f"{title}.bak"
+    block.download_file(title_bak)
+
+    assert os.path.isfile(title_bak)
+
+    os.remove(title_bak)
+
 
 def test_image_block(notion):
     block = notion.root_page.children.add_new(ImageBlock)
     assert_block_is_okay(**locals(), type="image")
+
     source = "https://raw.githubusercontent.com/jamalex/"
     source = source + "notion-py/master/ezgif-3-a935fdcb7415.gif"
+
     assert_block_attributes(block, source=source, caption="caption")
 
 
 def test_video_block(notion):
     block = notion.root_page.children.add_new(VideoBlock)
     assert_block_is_okay(**locals(), type="video")
+
     source = "https://streamable.com/8ud2kh"
+
     assert_block_attributes(block, source=source, caption="caption")