async-mega-py 2.0.3.dev0__tar.gz → 2.0.5.dev0__tar.gz

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: async-mega-py
-Version: 2.0.3.dev0
+Version: 2.0.5.dev0
 Summary: Python library for the Mega.nz and Transfer.it API
 Keywords: api,downloader,mega,mega.nz,transfer.it
 Author: NTFSvolume

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/pyproject.toml

@@ -34,7 +34,7 @@ license = "Apache-2.0"
 license-files = ["LICENSE"]
 readme = "README.md"
 requires-python = ">=3.11"
-version = "2.0.3.dev"
+version = "2.0.5.dev"
 
 [project.optional-dependencies]
 cli = [

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/src/mega/chunker.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import dataclasses
 import logging
 from collections.abc import Generator
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, NamedTuple
 
 from Crypto.Cipher import AES
 from Crypto.Util import Counter
@@ -17,9 +17,32 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+class ChunkBoundary(NamedTuple):
+    offset: int
+    size: int
+
+
+def get_chunks(size: int) -> Generator[ChunkBoundary]:
+    """
+    Yield chunk boundaries for Mega's MAC computation.
+
+    Chunk sizes double from 128 KiB (0x20000) up to 1 MiB (0x100000).
+    The last chunk may be smaller
+    """
+
+    offset = 0
+    current_size = init_size = 0x20000
+    while offset + current_size < size:
+        yield ChunkBoundary(offset, current_size)
+        offset += current_size
+        if current_size < 0x100000:
+            current_size += init_size
+    yield ChunkBoundary(offset, size - offset)
+
+
 @dataclasses.dataclass(slots=True)
 class MegaChunker:
-    """Decrypts/encrypts a flow of chunks using Mega's CBC algorithm"""
+    """Decrypts/encrypts a flow of chunks using Mega's custom CBC-MAC algorithm"""
 
     iv: tuple[int, int]
     key: tuple[int, int, int, int]

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/src/mega/client.py

@@ -22,7 +22,7 @@ from mega.crypto import (
 from mega.data_structures import AccountStats, Attributes, Crypto, FileInfo, Node, NodeID, NodeType, UserResponse
 from mega.download import DownloadResults
 from mega.filesystem import FileSystem
-from mega.utils import Site, throttled_gather
+from mega.utils import Site, async_map
 
 from .errors import MegaNzError, RequestError, ValidationError
 
@@ -155,7 +155,7 @@ class MegaNzClient(MegaCore):
             return await self.get_folder_link(fs[node.id])
 
     async def get_public_filesystem(self, public_handle: NodeID, public_key: str) -> FileSystem:
-        logger.info(f"Getting filesystem for {public_handle}...")
+        logger.info(f"Fetching filesystem information for {public_handle = }...")
         folder: GetNodesResponse = await self._api.post(
             {
                 "a": "f",
@@ -165,7 +165,8 @@ class MegaNzClient(MegaCore):
             },
             {"n": public_handle},
         )
-
+        nodes = folder["f"]
+        logger.info(f"Decrypting and building filesystem for {public_handle =} ({len(nodes)} nodes)...")
         nodes = await self._vault.deserialize_nodes(folder["f"], public_key)
         return await asyncio.to_thread(FileSystem.build, nodes)
 
@@ -207,7 +208,7 @@ class MegaNzClient(MegaCore):
         base_path = Path(output_dir or ".")
         folder_url = f"{_DOMAIN}/folder/{public_handle}#{public_key}"
 
-        async def worker(file: Node) -> tuple[NodeID, Path | Exception]:
+        async def download(file: Node) -> tuple[NodeID, Path | Exception]:
             web_url = folder_url + f"/file/{file.id}"
             output_path = base_path / fs.relative_path(file.id)
             try:
@@ -224,7 +225,7 @@ class MegaNzClient(MegaCore):
 
             return file.id, result
 
-        results = await throttled_gather(worker, fs.files_from(root_id))
+        results = await async_map(download, fs.files_from(root_id))
         return DownloadResults.split(dict(results))
 
     async def upload(self, file_path: str | PathLike[str], dest_node_id: NodeID | None = None) -> Node:

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/src/mega/core.py

@@ -22,7 +22,7 @@ from mega.crypto import (
 from mega.data_structures import Crypto, FileInfo, FileInfoSerialized, Node, NodeID
 from mega.errors import MegaNzError, RequestError, ValidationError
 from mega.filesystem import UserFileSystem
-from mega.utils import Site, random_u32int_array, transform_v1_url
+from mega.utils import Site, get_file_size, random_u32int_array, transform_v1_url
 from mega.vault import MegaVault
 
 if TYPE_CHECKING:
@@ -169,6 +169,7 @@ class MegaCore(AbstractApiClient):
         return FileInfo.parse(resp)
 
     async def _prepare_filesystem(self) -> UserFileSystem:
+        logger.info("Fetching users's filesystem information...")
         nodes_resp: GetNodesResponse = await self._api.post(
             {
                 "a": "f",
@@ -177,13 +178,15 @@ class MegaCore(AbstractApiClient):
             },
         )
 
+        nodes = nodes_resp["f"]
+        logger.info(f"Decrypting and building users's filesystem ({len(nodes)} nodes)...")
         self._vault.init_shared_keys(nodes_resp)
-        nodes = await self._vault.deserialize_nodes(nodes_resp["f"])
+        nodes = await self._vault.deserialize_nodes(nodes)
         return await asyncio.to_thread(UserFileSystem.build, nodes)
 
     async def _upload(self, file_path: str | PathLike[str], dest_node_id: NodeID) -> GetNodesResponse:
         file_path = Path(file_path)
-        file_size = file_path.stat().st_size
+        file_size = await asyncio.to_thread(get_file_size, file_path)
 
         with progress.new_task(file_path.name, file_size, "UP"):
             file_id, crypto = await upload.upload(self._api, file_path, file_size)

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/src/mega/crypto.py

@@ -7,14 +7,14 @@ import logging
 import math
 import struct
 import time
-from typing import TYPE_CHECKING, Any, NamedTuple
+from typing import TYPE_CHECKING, Any
 
 from Crypto.Cipher import AES
 from Crypto.Math.Numbers import Integer
 from Crypto.PublicKey import RSA
 
 if TYPE_CHECKING:
-    from collections.abc import Generator, Mapping, Sequence
+    from collections.abc import Mapping, Sequence
 
     from mega.data_structures import AttributesSerialized
 
@@ -24,11 +24,6 @@ CHUNK_BLOCK_LEN = 16  # Hexadecimal
 EMPTY_IV = b"\0" * CHUNK_BLOCK_LEN
 
 
-class ChunkBoundary(NamedTuple):
-    offset: int
-    size: int
-
-
 def pad_bytes(data: bytes | memoryview[int], length: int = CHUNK_BLOCK_LEN) -> bytes:
     if len(data) % length:
         padding = b"\0" * (length - len(data) % length)
@@ -147,18 +142,6 @@ def a32_to_base64(array: Sequence[int]) -> str:
     return b64_url_encode(a32_to_bytes(array))
 
 
-def get_chunks(size: int) -> Generator[ChunkBoundary]:
-    # generates a list of chunks (offset, chunk_size), where offset refers to the file initial position
-    offset = 0
-    current_size = init_size = 0x20000
-    while offset + current_size < size:
-        yield ChunkBoundary(offset, current_size)
-        offset += current_size
-        if current_size < 0x100000:
-            current_size += init_size
-    yield ChunkBoundary(offset, size - offset)
-
-
 def decrypt_rsa_key(private_key: bytes) -> RSA.RsaKey:
     # The private_key contains 4 MPI integers concatenated together.
     rsa_private_key = [0, 0, 0, 0]

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/src/mega/download.py

@@ -13,8 +13,7 @@ from types import MappingProxyType
 from typing import IO, TYPE_CHECKING, Final, Generic, Self, TypeVar
 
 from mega import progress
-from mega.chunker import MegaChunker
-from mega.crypto import get_chunks
+from mega.chunker import MegaChunker, get_chunks
 from mega.data_structures import NodeID
 
 if TYPE_CHECKING:
@@ -63,11 +62,11 @@ async def encrypted_stream(
 
         chunker = MegaChunker(iv, key, meta_mac)
         progress_hook = progress.current_hook.get()
-        async with _new_temp_download(output_path) as output:
+        async with _new_temp_download(output_path) as file_io:
             for _, chunk_size in get_chunks(file_size):
                 encrypted_chunk = await stream.readexactly(chunk_size)
                 chunk = chunker.read(encrypted_chunk)
-                output.write(chunk)
+                await asyncio.to_thread(file_io.write, chunk)
                 progress_hook(len(chunk))
 
             chunker.check_integrity()
@@ -81,9 +80,9 @@ async def stream(stream: aiohttp.StreamReader, output_path: Path) -> Path:
             raise FileExistsError(errno.EEXIST, output_path)
 
         progress_hook = progress.current_hook.get()
-        async with _new_temp_download(output_path) as output:
+        async with _new_temp_download(output_path) as file_io:
             async for chunk in stream.iter_chunked(_CHUNK_SIZE):
-                output.write(chunk)
+                await asyncio.to_thread(file_io.write, chunk)
                 progress_hook(len(chunk))
 
         return output_path
@@ -92,7 +91,7 @@ async def stream(stream: aiohttp.StreamReader, output_path: Path) -> Path:
 @contextlib.asynccontextmanager
 async def _new_temp_download(output_path: Path) -> AsyncGenerator[IO[bytes]]:
     # We need NamedTemporaryFile to not delete on file.close() but on context exit, which is not supported until python 3.12
-    temp_file = tempfile.NamedTemporaryFile(prefix="mega_py_", delete=False)
+    temp_file = await asyncio.to_thread(tempfile.NamedTemporaryFile, prefix="mega_py_", delete=False)
     logger.debug(f'Created temp file "{temp_file.name!s}" for "{output_path!s}"')
     try:
         yield temp_file

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/src/mega/progress.py

@@ -3,14 +3,16 @@ from __future__ import annotations
 import asyncio
 import contextlib
 from contextvars import ContextVar
-from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias
+from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias, TypeVar
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator
     from types import TracebackType
 
-    from rich.progress import Progress
+    from rich.progress import Progress, Task
+    from rich.text import Text
 
+    _T = TypeVar("_T")
     ProgressHook: TypeAlias = Callable[[float], None]
 
     class ProgressHookContext(Protocol):
@@ -33,12 +35,13 @@ current_hook: ContextVar[ProgressHook] = ContextVar("current_hook", default=lamb
 
 
 @contextlib.contextmanager
-def set_progress_factory(hook_factory: ProgressHookFactory) -> Generator[None]:
-    token = _PROGRESS_HOOK_FACTORY.set(hook_factory)
+def _enter_context(context_var: ContextVar[_T], value: _T) -> Generator[None]:
+    """Context manager for context vars"""
+    token = context_var.set(value)
     try:
         yield
     finally:
-        _PROGRESS_HOOK_FACTORY.reset(token)
+        context_var.reset(token)
 
 
 @contextlib.contextmanager
@@ -48,12 +51,8 @@ def new_task(description: str, total: float, kind: Literal["UP", "DOWN"]) -> Gen
         yield
         return
 
-    with factory(description, total, kind) as progress_hook:
-        token = current_hook.set(progress_hook)
-        try:
-            yield
-        finally:
-            current_hook.reset(token)
+    with factory(description, total, kind) as new_hook, _enter_context(current_hook, new_hook):
+        yield
 
 
 @contextlib.contextmanager
@@ -66,45 +65,64 @@ def new_progress() -> Generator[None]:
     def hook_factory(*args, **kwargs):
         return _new_rich_task(progress, *args, **kwargs)
 
-    with progress, set_progress_factory(hook_factory):
+    with (
+        progress,
+        _enter_context(_PROGRESS_HOOK_FACTORY, hook_factory),
+    ):
         yield
 
 
-def _truncate_desc(desc: str, length: int = 80, placeholder: str = "...") -> str:
-    if len(desc) <= length:
-        return desc
-
-    return f"{desc[: length - len(placeholder)]}{placeholder}"
-
-
 def _new_rich_progress() -> Progress | None:
     try:
+        from rich import get_console
         from rich.progress import (
             BarColumn,
             DownloadColumn,
             Progress,
             SpinnerColumn,
+            TextColumn,
             TimeRemainingColumn,
             TransferSpeedColumn,
         )
+        from rich.table import Column
     except ImportError:
         return None
 
-    else:
-        return Progress(
-            "[{task.fields[kind]}]",
-            SpinnerColumn(),
-            "{task.description}",
-            BarColumn(bar_width=None),
-            "[progress.percentage]{task.percentage:>6.2f}%",
-            "-",
-            DownloadColumn(),
-            "-",
-            TransferSpeedColumn(),
-            "-",
-            TimeRemainingColumn(compact=True, elapsed_when_finished=True),
-            transient=True,
-        )
+    console = get_console()
+
+    class AutoTruncatedTextColumn(TextColumn):
+        def render(self, task: Task) -> Text:
+            text = super().render(task)
+            width = console.width
+            available_witdh = min((width * 60 // 100), (width - 65))
+            desc_limit = max(available_witdh, 8)
+            text.truncate(desc_limit, overflow="ellipsis")
+            return text
+
+    return Progress(
+        "[{task.fields[kind]}]",
+        SpinnerColumn(),
+        AutoTruncatedTextColumn("{task.description}"),
+        BarColumn(
+            bar_width=None,
+        ),
+        "[progress.percentage]{task.percentage:>6.1f}%",
+        "•",
+        DownloadColumn(
+            table_column=Column(justify="right", no_wrap=True),
+        ),
+        "•",
+        TransferSpeedColumn(table_column=Column(justify="right", no_wrap=True)),
+        "•",
+        TimeRemainingColumn(
+            compact=True,
+            elapsed_when_finished=True,
+            table_column=Column(justify="right", no_wrap=True),
+        ),
+        transient=True,
+        console=console,
+        expand=True,
+    )
 
 
 @contextlib.contextmanager
@@ -114,7 +132,7 @@ def _new_rich_task(
     total: float,
     kind: Literal["UP", "DOWN"],
 ) -> Generator[ProgressHook]:
-    task_id = progress.add_task(_truncate_desc(description), total=total, kind=kind)
+    task_id = progress.add_task(description, total=total, kind=kind)
 
     def progress_hook(advance: float) -> None:
         progress.advance(task_id, advance)

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/src/mega/transfer_it.py

@@ -15,7 +15,7 @@ from mega.crypto import b64_to_a32, b64_url_decode, decrypt_attr
 from mega.data_structures import Attributes, Crypto, Node, NodeID, NodeType
 from mega.download import DownloadResults
 from mega.filesystem import FileSystem
-from mega.utils import Site, throttled_gather
+from mega.utils import Site, async_map
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -47,6 +47,7 @@ class TransferItClient(AbstractApiClient):
         self._api = TransferItAPI(session)
 
     async def get_filesystem(self, transfer_id: TransferID) -> FileSystem:
+        logger.info(f"Fetching filesystem information for {transfer_id = }...")
         folder: GetNodesResponse = await self._api.post(
             {
                 "a": "f",
@@ -56,7 +57,9 @@ class TransferItClient(AbstractApiClient):
             },
             {"x": transfer_id},
         )
-        return await asyncio.to_thread(self._deserialize_nodes, folder["f"])
+        nodes = folder["f"]
+        logger.info(f"Decrypting and building filesystem for {transfer_id = } ({len(nodes)} nodes)...")
+        return await asyncio.to_thread(self._deserialize_nodes, nodes)
 
     @staticmethod
     def parse_url(url: str | yarl.URL) -> TransferID:
@@ -107,7 +110,7 @@ class TransferItClient(AbstractApiClient):
         base_path = Path(output_dir or ".") / f"transfer.it ({transfer_id})"
         folder_url = f"https://transfer.it/t/{transfer_id}"
 
-        async def worker(file: Node) -> tuple[NodeID, Path | Exception]:
+        async def download(file: Node) -> tuple[NodeID, Path | Exception]:
             web_url = folder_url + f"#{file.id}"
             output_path = base_path / fs.relative_path(file.id)
             dl_link = self.create_download_url(transfer_id, file)
@@ -123,7 +126,7 @@ class TransferItClient(AbstractApiClient):
 
             return file.id, result
 
-        results = await throttled_gather(worker, fs.files_from(root_id))
+        results = await async_map(download, fs.files_from(root_id))
         return DownloadResults.split(dict(results))
 
     async def _download_file(self, dl_link: str, output_path: str | PathLike[str]) -> Path:

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/src/mega/upload.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
+import asyncio
 import logging
-from typing import IO, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 from mega import progress
-from mega.chunker import MegaChunker
-from mega.crypto import a32_to_base64, b64_url_encode, encrypt_attr, encrypt_key, get_chunks
-from mega.data_structures import Crypto
+from mega.chunker import MegaChunker, get_chunks
+from mega.crypto import a32_to_base64, b64_url_encode, encrypt_attr, encrypt_key
+from mega.data_structures import ByteSize, Crypto
 from mega.utils import random_u32int_array
 
 if TYPE_CHECKING:
@@ -24,24 +25,23 @@ async def _request_upload_url(api: MegaAPI, file_size: int) -> str:
 
 
 async def upload(api: MegaAPI, file_path: Path, file_size: int) -> tuple[str, Crypto]:
-    with file_path.open("rb") as input_file:
-        random_array = random_u32int_array(6)
-        key, iv = random_array[:4], random_array[4:]
+    random_array = random_u32int_array(6)
+    key, iv = random_array[:4], random_array[4:]
 
-        if file_size == 0:
-            upload_url = await _request_upload_url(api, file_size)
-            file_handle = await api.upload_chunk(upload_url, 0, b"")
-            meta_mac = 0, 0
-            return file_handle, Crypto.compose(key, iv, meta_mac)
+    if file_size == 0:
+        upload_url = await _request_upload_url(api, file_size)
+        file_handle = await api.upload_chunk(upload_url, 0, b"")
+        meta_mac = 0, 0
+        return file_handle, Crypto.compose(key, iv, meta_mac)
 
-        chunker = MegaChunker(iv, key)  # pyright: ignore[reportArgumentType]
-        return await _upload_chunks(api, chunker, input_file, file_size)
+    chunker = MegaChunker(iv, key)  # pyright: ignore[reportArgumentType]
+    return await _upload_chunks(api, chunker, file_path, file_size)
 
 
 async def _upload_chunks(
     api: MegaAPI,
     chunker: MegaChunker,
-    input_file: IO[bytes],
+    file_path: Path,
     file_size: int,
 ) -> tuple[str, Crypto]:
     upload_progress = 0
@@ -49,13 +49,18 @@ async def _upload_chunks(
     upload_url = await _request_upload_url(api, file_size)
     progress_hook = progress.current_hook.get()
 
-    for offset, size in get_chunks(file_size):
-        chunk = chunker.read(input_file.read(size))
-        file_handle = await api.upload_chunk(upload_url, offset, chunk)
-        logger.info(f"{upload_progress} of {file_size} uploaded ({upload_progress / file_size:0.1f}%)")
-        real_size = len(chunk)
-        upload_progress += real_size
-        progress_hook(real_size)
+    file_size = ByteSize(file_size)
+    total = file_size.human_readable()
+    with await asyncio.to_thread(file_path.open, "rb") as input_file:
+        for offset, size in get_chunks(file_size):
+            chunk = chunker.read(await asyncio.to_thread(input_file.read, size))
+            file_handle = await api.upload_chunk(upload_url, offset, chunk)
+            human_progress = ByteSize(upload_progress).human_readable()
+            ratio = (upload_progress / file_size) * 100
+            logger.debug(f'{human_progress}/{total} uploaded ({ratio:0.1f}%) for "{file_path!s}"')
+            real_size = len(chunk)
+            upload_progress += real_size
+            progress_hook(real_size)
 
     assert file_handle
     return file_handle, Crypto.compose(chunker.key, chunker.iv, chunker.compute_meta_mac())

{async_mega_py-2.0.3.dev0 → async_mega_py-2.0.5.dev0}/src/mega/utils.py

@@ -2,16 +2,19 @@ from __future__ import annotations
 
 import asyncio
 import datetime
+import errno
 import logging
 import random
 import string
 from enum import Enum
+from stat import S_ISREG
 from typing import TYPE_CHECKING, Literal, TypeVar, overload
 
 import yarl
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable, Iterable, Sequence
+    from pathlib import Path
 
     _T1 = TypeVar("_T1")
     _T2 = TypeVar("_T2")
@@ -76,8 +79,20 @@ def transform_v1_url(url: yarl.URL) -> yarl.URL:
     return url
 
 
+def get_file_size(file_path: Path) -> int:
+    try:
+        stat = file_path.stat()
+    except (OSError, ValueError):
+        raise FileNotFoundError(errno.ENOENT, str(file_path)) from None
+
+    if not S_ISREG(stat.st_mode):
+        raise IsADirectoryError(errno.EISDIR, str(file_path))
+
+    return stat.st_size
+
+
 @overload
-async def throttled_gather(
+async def async_map(
     coro_factory: Callable[[_T1], Awaitable[_T2]],
     values: Iterable[_T1],
     *,
@@ -87,7 +102,7 @@ async def throttled_gather(
 
 
 @overload
-async def throttled_gather(
+async def async_map(
     coro_factory: Callable[[_T1], Awaitable[_T2]],
     values: Iterable[_T1],
     *,
@@ -96,7 +111,7 @@ async def throttled_gather(
 ) -> list[_T2]: ...
 
 
-async def throttled_gather(
+async def async_map(
     coro_factory: Callable[[_T1], Awaitable[_T2]],
     values: Iterable[_T1],
     *,