async-mega-py 2.0.1.dev0__tar.gz → 2.0.4.dev0__tar.gz

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: async-mega-py
-Version: 2.0.1.dev0
+Version: 2.0.4.dev0
 Summary: Python library for the Mega.nz and Transfer.it API
 Keywords: api,downloader,mega,mega.nz,transfer.it
 Author: NTFSvolume
@@ -167,8 +167,8 @@ await mega.download_public_file(public_handle, public_key, output_dir)
 
 # Download a public folder
 url = "https://mega.co.nz/#F!utYjgSTQ!OM4U3V5v_W4N5edSo0wolg1D5H0fwSrLD3oLnLuS9pc"
-public_handle, public_key = mega.parse_folder_url(url)
-success, fails = await mega.download_public_folder(public_handle, public_key, output_dir)
+public_handle, public_key, selected_node = mega.parse_folder_url(url)
+success, fails = await mega.download_public_folder(public_handle, public_key, output_dir, selected_node)
 print(f"Download of '{url!s}' finished. Successful downloads {len(success)}, failed {len(fails)}")
 
 # Import a file from URL
@@ -351,21 +351,26 @@ uv tool install async-mega-py[cli]
 mega-py --help
 ```
 
-```sh
- Usage: mega-py [OPTIONS] COMMAND [ARGS]...  
+```powershell
+Usage: mega-py [OPTIONS] COMMAND [ARGS]...  
+
+ CLI app for the Mega.nz and Transfer.it. Set MEGA_NZ_EMAIL and MEGA_NZ_PASSWORD  
+ enviroment variables to use them as credentials for Mega  
 
 ╭─ Options ──────────────────────────────────────────────────────────────────────╮
-│ --help          Show this message and exit.                                    │
+│ --verbose  -v               Increase verbosity (-v shows debug logs,  -vv      │
+│                             shows HTTP traffic)                                │
+│ --help                      Show this message and exit.                        │
 ╰────────────────────────────────────────────────────────────────────────────────╯
 ╭─ Commands ─────────────────────────────────────────────────────────────────────╮
 │ download   Download a public file or folder by its URL (transfer.it / mega.nz) │
 │ dump       Dump a copy of your filesystem to disk                              │
 │ stats      Show account stats                                                  │
 │ upload     Upload a file to your account                                       │
-╰────────────────────────────────────────────
+╰────────────────────────────────────────────────────────────────────────────────╯
 ```
 
 > [!TIP]
-> The CLI app does _not_ accept login credentials, but you can still use your account by setting up the `MEGA_NZ_EMAIL` and `MEGA_NZ_PASSWORD` enviroment variables
+> The CLI app does *not* accept login credentials, but you can still use your account by setting up the `MEGA_NZ_EMAIL` and `MEGA_NZ_PASSWORD` enviroment variables
 >
 > It will also read them from an `.env` file (if found)

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/README.md

@@ -131,8 +131,8 @@ await mega.download_public_file(public_handle, public_key, output_dir)
 
 # Download a public folder
 url = "https://mega.co.nz/#F!utYjgSTQ!OM4U3V5v_W4N5edSo0wolg1D5H0fwSrLD3oLnLuS9pc"
-public_handle, public_key = mega.parse_folder_url(url)
-success, fails = await mega.download_public_folder(public_handle, public_key, output_dir)
+public_handle, public_key, selected_node = mega.parse_folder_url(url)
+success, fails = await mega.download_public_folder(public_handle, public_key, output_dir, selected_node)
 print(f"Download of '{url!s}' finished. Successful downloads {len(success)}, failed {len(fails)}")
 
 # Import a file from URL
@@ -315,21 +315,26 @@ uv tool install async-mega-py[cli]
 mega-py --help
 ```
 
-```sh
- Usage: mega-py [OPTIONS] COMMAND [ARGS]...  
+```powershell
+Usage: mega-py [OPTIONS] COMMAND [ARGS]...  
+
+ CLI app for the Mega.nz and Transfer.it. Set MEGA_NZ_EMAIL and MEGA_NZ_PASSWORD  
+ enviroment variables to use them as credentials for Mega  
 
 ╭─ Options ──────────────────────────────────────────────────────────────────────╮
-│ --help          Show this message and exit.                                    │
+│ --verbose  -v               Increase verbosity (-v shows debug logs,  -vv      │
+│                             shows HTTP traffic)                                │
+│ --help                      Show this message and exit.                        │
 ╰────────────────────────────────────────────────────────────────────────────────╯
 ╭─ Commands ─────────────────────────────────────────────────────────────────────╮
 │ download   Download a public file or folder by its URL (transfer.it / mega.nz) │
 │ dump       Dump a copy of your filesystem to disk                              │
 │ stats      Show account stats                                                  │
 │ upload     Upload a file to your account                                       │
-╰────────────────────────────────────────────
+╰────────────────────────────────────────────────────────────────────────────────╯
 ```
 
 > [!TIP]
-> The CLI app does _not_ accept login credentials, but you can still use your account by setting up the `MEGA_NZ_EMAIL` and `MEGA_NZ_PASSWORD` enviroment variables
+> The CLI app does *not* accept login credentials, but you can still use your account by setting up the `MEGA_NZ_EMAIL` and `MEGA_NZ_PASSWORD` enviroment variables
 >
 > It will also read them from an `.env` file (if found)

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/pyproject.toml

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

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/src/mega/api.py

@@ -5,6 +5,7 @@ import contextlib
 import dataclasses
 import logging
 from collections.abc import Mapping, Sequence
+from contextvars import ContextVar
 from functools import wraps
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, ParamSpec, Self, TypeVar
@@ -28,6 +29,7 @@ _UA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:137.0) Gecko/20100101 Firefo
 _DEFAULT_HEADERS: MappingProxyType[str, str] = MappingProxyType({"User-Agent": _UA})
 
 
+LOG_HTTP_TRAFFIC: ContextVar[bool] = ContextVar("LOG_HTTP_TRAFFIC", default=False)
 logger = logging.getLogger(__name__)
 
 
@@ -84,6 +86,7 @@ class MegaAPI:
         self.__session = session
         self._auto_close_session = session is None
         self._rate_limiter = AsyncLimiter(100, 60)
+        logger.disabled = not LOG_HTTP_TRAFFIC.get()
 
     @property
     def entrypoint(self) -> yarl.URL:

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/src/mega/cli/__init__.py

@@ -10,6 +10,7 @@ import typer
 import yarl
 
 from mega import __version__, env, progress
+from mega.api import LOG_HTTP_TRAFFIC
 from mega.cli.app import CLIApp
 from mega.client import MegaNzClient
 from mega.transfer_it import TransferItClient
@@ -22,8 +23,22 @@ if TYPE_CHECKING:
 logger = logging.getLogger("mega")
 
 
-def verbose(verbose: Annotated[bool, typer.Option("--verbose")] = False) -> None:
-    setup_logger(logging.DEBUG if verbose else logging.INFO)
+def verbose(
+    verbose: Annotated[
+        int,
+        typer.Option(
+            "-v",
+            "--verbose",
+            count=True,
+            help="Increase verbosity (-v shows debug logs,  -vv shows HTTP traffic)",
+        ),
+    ] = 0,
+) -> None:
+    if verbose > 1:
+        LOG_HTTP_TRAFFIC.set(True)
+
+    level = logging.DEBUG if verbose else logging.INFO
+    setup_logger(level)
 
 
 app = CLIApp(
@@ -51,8 +66,10 @@ async def transfer_it(url: str, output_dir: Path) -> None:
         with progress.new_progress():
             transfer_id = client.parse_url(url)
             logger.info(f"Downloading '{url}'")
-            success, fails = await client.download_transfer(transfer_id, output_dir / transfer_id)
-            logger.info(f"Download of '{url}' finished. Successful downloads {len(success)}, failed {len(fails)}")
+            results = await client.download_transfer(transfer_id, output_dir)
+            logger.info(
+                f"Download of '{url}' finished. Successful = {len(results.success)}, failed = {len(results.fails)}"
+            )
 
 
 @app.command()
@@ -112,10 +129,10 @@ async def download_file(mega: MegaNzClient, url: str, output: Path) -> None:
 
 
 async def download_folder(mega: MegaNzClient, url: str, output: Path) -> None:
-    public_handle, public_key = mega.parse_folder_url(url)
+    public_handle, public_key, root_node = mega.parse_folder_url(url)
     logger.info(f"Downloading {url}")
-    success, fails = await mega.download_public_folder(public_handle, public_key, output)
-    logger.info(f"Download of {url} finished. Successful downloads {len(success)}, failed {len(fails)}")
+    results = await mega.download_public_folder(public_handle, public_key, output, root_node)
+    logger.info(f"Download of '{url}' finished. Successful = {len(results.success)}, failed = {len(results.fails)}")
 
 
 def main() -> None:

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/src/mega/client.py

@@ -6,6 +6,8 @@ import logging
 from pathlib import Path, PurePosixPath
 from typing import TYPE_CHECKING, Any
 
+import aiohttp
+
 from mega import progress
 from mega.core import MegaCore
 from mega.crypto import (
@@ -18,8 +20,9 @@ from mega.crypto import (
     encrypt_key,
 )
 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 throttled_gather
+from mega.utils import Site, async_map
 
 from .errors import MegaNzError, RequestError, ValidationError
 
@@ -33,6 +36,9 @@ __all__ = ["MegaNzClient"]
 logger = logging.getLogger(__name__)
 
 
+_DOMAIN = Site.MEGA.value
+
+
 class MegaNzClient(MegaCore):
     """Interface with all the public methods of the API"""
 
@@ -95,7 +101,7 @@ class MegaNzClient(MegaCore):
 
         public_handle = await self._get_public_handle(file.id)
         public_key = a32_to_base64(file._crypto.full_key)
-        return f"{self._primary_url}/file/{public_handle}#{public_key}"
+        return f"{_DOMAIN}/file/{public_handle}#{public_key}"
 
     async def get_folder_link(self, folder: Node) -> str:
         if folder.type is not NodeType.FOLDER:
@@ -105,7 +111,7 @@ class MegaNzClient(MegaCore):
             raise RequestError("")
         public_handle = await self._get_public_handle(folder.id)
         public_key = a32_to_base64(folder._crypto.share_key)
-        return f"{self._primary_url}/#F!{public_handle}!{public_key}"
+        return f"{_DOMAIN}/folder/{public_handle}#{public_key}"
 
     async def get_id_from_public_handle(self, public_handle: NodeID) -> str:
         resp: GetNodesResponse = await self._api.post(
@@ -149,6 +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}...")
         folder: GetNodesResponse = await self._api.post(
             {
                 "a": "f",
@@ -165,11 +172,8 @@ class MegaNzClient(MegaCore):
     async def download(self, node: Node, output_dir: str | PathLike[str] | None = None) -> Path:
         """Download a file by it's file object."""
         file_info = await self._request_file_info(node.id)
-        return await self._download_file(
-            file_info,
-            node._crypto,
-            output_folder=output_dir,
-        )
+        output_path = Path(output_dir or ".") / node.attributes.name
+        return await self._download_file(file_info, node._crypto, output_path)
 
     async def download_public_file(
         self,
@@ -180,18 +184,18 @@ class MegaNzClient(MegaCore):
         full_key = b64_to_a32(public_key)
         crypto = Crypto.decompose(full_key)
         file_info = await self._request_file_info(public_handle, is_public=True)
-        return await self._download_file(
-            file_info,
-            crypto,
-            output_dir,
-        )
+        attrs = decrypt_attr(b64_url_decode(file_info._at), crypto.key)
+        output_name = Attributes.parse(attrs).name
+        output_path = Path(output_dir or ".") / output_name
+        return await self._download_file(file_info, crypto, output_path)
 
     async def download_public_folder(
         self,
         public_handle: NodeID,
         public_key: str,
         output_dir: str | PathLike[str] | None = None,
-    ) -> tuple[list[Path], list[Exception]]:
+        root_id: NodeID | None = None,
+    ) -> DownloadResults:
         """Recursively download all files from a public folder, preserving its internal directory structure.
 
         Returns:
@@ -201,30 +205,27 @@ class MegaNzClient(MegaCore):
         fs = await self.get_public_filesystem(public_handle, public_key)
 
         base_path = Path(output_dir or ".")
-        folder_url = f"{self._primary_url}/folder/{public_handle}#{public_key}"
+        folder_url = f"{_DOMAIN}/folder/{public_handle}#{public_key}"
 
-        async def worker(file: Node, path: PurePosixPath) -> Path:
+        async def download(file: Node) -> tuple[NodeID, Path | Exception]:
             web_url = folder_url + f"/file/{file.id}"
-            output_folder = base_path / path.parent
+            output_path = base_path / fs.relative_path(file.id)
             try:
                 file_info = await self._request_file_info(file.id, public_handle)
-                return await self._download_file(file_info, file._crypto, output_folder, path.name)
-            except Exception as exc:
-                logger.error(f'Unable to download {web_url} to "{output_folder}" ({exc})')
-                raise
+                result = await self._download_file(file_info, file._crypto, output_path)
 
-        def make_coros():
-            for file in fs.files:
-                path = fs.relative_path(file.id)
-                yield (worker(file, path))
+            except Exception as exc:
+                if isinstance(exc, aiohttp.ClientResponseError):
+                    msg = f"[{exc.status}] {exc.message}"
+                else:
+                    msg = f"({type(exc).__name__})"
+                logger.error(f'Unable to download {web_url} to "{output_path}" {msg}')
+                result = exc
 
-        results = await throttled_gather(make_coros(), return_exceptions=True)
-        success: list[Path] = []
-        fails: list[Exception] = [
-            result for result in results if isinstance(result, Exception) or (success.append(result) and False)
-        ]
+            return file.id, result
 
-        return success, fails
+        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:
         if not dest_node_id:

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/src/mega/core.py

@@ -5,7 +5,7 @@ import logging
 import random
 import re
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 import yarl
 from Crypto.Cipher import AES
@@ -15,13 +15,11 @@ from mega.api import AbstractApiClient
 from mega.crypto import (
     a32_to_base64,
     a32_to_bytes,
-    b64_url_decode,
     b64_url_encode,
-    decrypt_attr,
     encrypt_attr,
     encrypt_key,
 )
-from mega.data_structures import Attributes, Crypto, FileInfo, FileInfoSerialized, Node, NodeID
+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
@@ -40,15 +38,18 @@ logger = logging.getLogger(__name__)
 
 class ParsedPublicURL(NamedTuple):
     is_folder: bool
-    public_handle: str
+    public_handle: NodeID
     public_key: str
-    inner_folder_id: str | None = None
-    inner_file_id: str | None = None
+    selected_folder: NodeID | None = None
+    selected_file: NodeID | None = None
+
+    @property
+    def selected_node(self) -> NodeID | None:
+        return self.selected_folder or self.selected_file
 
 
 class MegaCore(AbstractApiClient):
     __slots__ = ("_filesystem", "_lock", "_vault")
-    _primary_url: ClassVar[str] = "https://mega.nz"
 
     def __init__(self, session: aiohttp.ClientSession | None = None) -> None:
         super().__init__(session)
@@ -89,18 +90,18 @@ class MegaCore(AbstractApiClient):
         logger.info("Login complete!")
 
     @classmethod
-    def parse_file_url(cls, url: str | yarl.URL) -> tuple[str, str]:
+    def parse_file_url(cls, url: str | yarl.URL) -> tuple[NodeID, str]:
         result = cls.parse_url(url)
         if result.is_folder:
             raise ValueError("This is a folder URL: {url}")
         return result.public_handle, result.public_key
 
     @classmethod
-    def parse_folder_url(cls, url: str | yarl.URL) -> tuple[str, str]:
+    def parse_folder_url(cls, url: str | yarl.URL) -> tuple[NodeID, str, NodeID | None]:
         result = cls.parse_url(url)
         if not result.is_folder:
             raise ValueError("This is a file URL: {url}")
-        return result.public_handle, result.public_key
+        return result.public_handle, result.public_key, result.selected_node
 
     @staticmethod
     def parse_url(url: str | yarl.URL) -> ParsedPublicURL:
@@ -200,8 +201,7 @@ class MegaCore(AbstractApiClient):
         self,
         file_info: FileInfo,
         crypto: Crypto,
-        output_folder: str | PathLike[str] | None = None,
-        output_name: str | None = None,
+        output_folder: str | PathLike[str],
     ) -> Path:
         # Seems to happens sometime... When this occurs, files are
         # inaccessible also in the official web app.
@@ -209,12 +209,7 @@ class MegaCore(AbstractApiClient):
         if not file_info.url:
             raise RequestError("File not accessible anymore")
 
-        if not output_name:
-            attrs = decrypt_attr(b64_url_decode(file_info._at), crypto.key)
-            output_name = Attributes.parse(attrs).name
-
-        output_path = Path(output_folder or Path()) / output_name
-
+        output_path = Path(output_folder)
         async with self._api.get(file_info.url) as response:
             with progress.new_task(output_path.name, file_info.size, "DOWN"):
                 return await download.encrypted_stream(

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/src/mega/data_structures.py

@@ -222,6 +222,14 @@ class Node(_DictDumper):
     _a: str
     _crypto: Crypto
 
+    @property
+    def is_file(self) -> bool:
+        return self.type is NodeType.FILE
+
+    @property
+    def is_folder(self) -> bool:
+        return self.type is NodeType.FOLDER
+
     @classmethod
     def parse(cls, node: NodeSerialized) -> Node:
         owner = node.get("u", "")

async_mega_py-2.0.4.dev0/src/mega/download.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import errno
+import logging
+import shutil
+import tempfile
+import weakref
+from collections.abc import Iterator, Mapping
+from pathlib import Path
+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.data_structures import NodeID
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+    import aiohttp
+
+_T = TypeVar("_T")
+logger = logging.getLogger(__name__)
+
+
+class WeakAsyncLocks(Generic[_T]):
+    """A WeakValueDictionary wrapper for asyncio.Locks.
+
+    Unused locks are automatically garbage collected. When trying to retrieve a
+    lock that does not exists, a new lock will be created.
+    """
+
+    __slots__ = ("__locks",)
+
+    def __init__(self) -> None:
+        self.__locks: Final = weakref.WeakValueDictionary[_T, asyncio.Lock]()
+
+    def __getitem__(self, key: _T, /) -> asyncio.Lock:
+        lock = self.__locks.get(key)
+        if lock is None:
+            self.__locks[key] = lock = asyncio.Lock()
+        return lock
+
+
+_LOCKS: WeakAsyncLocks[Path] = WeakAsyncLocks()
+_CHUNK_SIZE = 1024 * 1024 * 5  # 5MB
+
+
+async def encrypted_stream(
+    stream: aiohttp.StreamReader,
+    output_path: Path,
+    file_size: int,
+    iv: tuple[int, int],
+    meta_mac: tuple[int, int],
+    key: tuple[int, int, int, int],
+) -> Path:
+    async with _LOCKS[output_path]:
+        if await asyncio.to_thread(output_path.exists):
+            raise FileExistsError(errno.EEXIST, output_path)
+
+        chunker = MegaChunker(iv, key, meta_mac)
+        progress_hook = progress.current_hook.get()
+        async with _new_temp_download(output_path) as output:
+            for _, chunk_size in get_chunks(file_size):
+                encrypted_chunk = await stream.readexactly(chunk_size)
+                chunk = chunker.read(encrypted_chunk)
+                output.write(chunk)
+                progress_hook(len(chunk))
+
+            chunker.check_integrity()
+
+        return output_path
+
+
+async def stream(stream: aiohttp.StreamReader, output_path: Path) -> Path:
+    async with _LOCKS[output_path]:
+        if await asyncio.to_thread(output_path.exists):
+            raise FileExistsError(errno.EEXIST, output_path)
+
+        progress_hook = progress.current_hook.get()
+        async with _new_temp_download(output_path) as output:
+            async for chunk in stream.iter_chunked(_CHUNK_SIZE):
+                output.write(chunk)
+                progress_hook(len(chunk))
+
+        return output_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)
+    logger.debug(f'Created temp file "{temp_file.name!s}" for "{output_path!s}"')
+    try:
+        yield temp_file
+
+        def move():
+            temp_file.close()
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            shutil.move(temp_file.name, output_path)
+            logger.debug(f'Moved temp file "{temp_file.name!s}" to "{output_path!s}"')
+
+        await asyncio.to_thread(move)
+
+    finally:
+
+        def delete():
+            if not temp_file.closed:
+                temp_file.close()
+            Path(temp_file.name).unlink(missing_ok=True)
+
+        await asyncio.to_thread(delete)
+
+
+class DownloadResults(Mapping[NodeID, Path | Exception]):
+    success: MappingProxyType[NodeID, Path]
+    fails: MappingProxyType[NodeID, Exception]
+
+    def __init__(self, success: Mapping[NodeID, Path], fails: Mapping[NodeID, Exception]) -> None:
+        if not success.keys().isdisjoint(fails.keys()):
+            raise ValueError("A NodeID cannot be in both success and fails")
+        self.success = MappingProxyType(success)
+        self.fails = MappingProxyType(fails)
+
+    def __getitem__(self, value: NodeID) -> Path | Exception:
+        try:
+            return self.success[value]
+        except KeyError:
+            return self.fails[value]
+
+    def __iter__(self) -> Iterator[NodeID]:
+        yield from self.success
+        yield from self.fails
+
+    def __len__(self) -> int:
+        return len(self.success) + len(self.fails)
+
+    @classmethod
+    def split(cls, results: Mapping[NodeID, Path | Exception]) -> Self:
+        success: dict[NodeID, Path] = {}
+        fails: dict[NodeID, Exception] = {}
+
+        for node_id, result in results.items():
+            if isinstance(result, Exception):
+                fails[node_id] = result
+            else:
+                success[node_id] = result
+
+        return cls(success, fails)

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/src/mega/filesystem.py

@@ -7,7 +7,6 @@ from __future__ import annotations
 
 import dataclasses
 import errno
-from collections.abc import Generator
 from pathlib import PurePosixPath
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, NamedTuple, Self
@@ -34,7 +33,7 @@ def _resolve_paths(walker: _NodeWalker, *roots: Node) -> Generator[_NodeLookup]:
             node_path = current_path / node.attributes.name
             yield node.id, node_path
 
-            if node.type is not NodeType.FILE:
+            if not node.is_file:
                 yield from walk(node.id, node_path)
 
     for root in roots:
@@ -136,7 +135,7 @@ class SimpleFileSystem(_NodeWalker, _DictDumper):
                 case NodeType.TRASH:
                     trash_bin = node
                 case _:
-                    raise RuntimeError
+                    raise RuntimeError  # pyright: ignore[reportUnreachable]
 
         return cls(
             root=root,
@@ -234,14 +233,14 @@ class FileSystem(SimpleFileSystem):
     def files(self) -> Iterable[Node]:
         """All files that are NOT deleted (recursive)"""
         for node in self:
-            if node.type is NodeType.FILE and not self._deleted:
+            if node.is_file and node.id not in self._deleted:
                 yield node
 
     @property
     def folders(self) -> Iterable[Node]:
         """All folders that are NOT deleted (recursive)"""
         for node in self:
-            if node.type is NodeType.FOLDER and node.id not in self._deleted:
+            if node.is_folder and node.id not in self._deleted:
                 yield node
 
     def dirmap(self, node_id: str, *, recursive: bool = False) -> dict[NodeID, PurePosixPath]:
@@ -305,6 +304,27 @@ class FileSystem(SimpleFileSystem):
             assert nodes
             return self[nodes[0]]
 
+    def files_from(self, node_id: NodeID | None) -> Iterable[Node]:
+        """
+        Yield every file that is reachable from `node_id`.
+
+        - If `node_id` is `None`: yield all non deleted files on the file system.
+        - If `node_id` points to a file: yield only that file.
+        - Any other case: yield all files within that node (recursively).
+
+        """
+        if not node_id:
+            yield from self.files
+            return
+
+        root = self[node_id]
+        if root.is_file:
+            yield root
+        else:
+            for child in self.iterdir(root.id, recursive=True):
+                if child.is_file:
+                    yield child
+
     def dump(self, *, simple: bool = False) -> dict[str, Any]:
         """Get a JSONable dict representation of this object"""
         dump = super(FileSystem, self).dump()

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.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(_truncate_desc(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 = 30, 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

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/src/mega/transfer_it.py

@@ -3,24 +3,24 @@ from __future__ import annotations
 import asyncio
 import dataclasses
 import logging
-from pathlib import Path, PurePosixPath
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias
 
+import aiohttp
 import yarl
 
 from mega import download, progress
 from mega.api import AbstractApiClient, MegaAPI
 from mega.crypto import b64_to_a32, b64_url_decode, decrypt_attr
-from mega.data_structures import Attributes, Crypto, Node, NodeType
+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
     from os import PathLike
 
-    import aiohttp
-
     from mega.data_structures import GetNodesResponse, NodeSerialized
 
 
@@ -94,7 +94,8 @@ class TransferItClient(AbstractApiClient):
         self,
         transfer_id: TransferID,
         output_dir: str | PathLike[str] | None = None,
-    ) -> tuple[list[Path], list[Exception]]:
+        root_id: NodeID | None = None,
+    ) -> DownloadResults:
         """Recursively download all files from a transfer, preserving its internal directory structure.
 
         Returns:
@@ -103,42 +104,31 @@ class TransferItClient(AbstractApiClient):
         """
         fs = await self.get_filesystem(transfer_id)
 
-        base_path = Path(output_dir or ".")
+        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, path: PurePosixPath) -> Path:
+        async def download(file: Node) -> tuple[NodeID, Path | Exception]:
             web_url = folder_url + f"#{file.id}"
-            output_folder = base_path / path.parent
+            output_path = base_path / fs.relative_path(file.id)
             dl_link = self.create_download_url(transfer_id, file)
             try:
-                return await self._download_file(dl_link, output_folder, path.name)
+                result = await self._download_file(dl_link, output_path)
             except Exception as exc:
-                logger.error(f'Unable to download {web_url} to "{output_folder}" ({exc})')
-                raise
-
-        def make_coros():
-            for file in fs.files:
-                path = fs.relative_path(file.id)
-                yield (worker(file, path))
+                if isinstance(exc, aiohttp.ClientResponseError):
+                    msg = f"[{exc.status}] {exc.message}"
+                else:
+                    msg = f"({type(exc).__name__})"
+                logger.error(f'Unable to download {web_url} to "{output_path}" {msg}')
+                result = exc
 
-        results = await throttled_gather(make_coros(), return_exceptions=True)
-        success: list[Path] = []
-        fails: list[Exception] = [
-            result for result in results if isinstance(result, Exception) or (success.append(result) and False)
-        ]
+            return file.id, result
 
-        return success, fails
-
-    async def _download_file(
-        self,
-        dl_link: str,
-        output_folder: str | PathLike[str] | None = None,
-        output_name: str | None = None,
-    ) -> Path:
-        name = output_name or yarl.URL(dl_link).query["fn"]
-        output_path = Path(output_folder or Path()) / name
+        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:
+        output_path = Path(output_path)
         async with self._api.get(dl_link, headers={"Referer": "https://transfer.it/"}) as response:
             size = int(response.headers["Content-Length"])
-            with progress.new_task(name, size, "DOWN"):
+            with progress.new_task(output_path.name, size, "DOWN"):
                 return await download.stream(response.content, output_path)

{async_mega_py-2.0.1.dev0 → async_mega_py-2.0.4.dev0}/src/mega/utils.py

@@ -11,9 +11,10 @@ from typing import TYPE_CHECKING, Literal, TypeVar, overload
 import yarl
 
 if TYPE_CHECKING:
-    from collections.abc import Awaitable, Iterable, Sequence
+    from collections.abc import Awaitable, Callable, Iterable, Sequence
 
-    _T = TypeVar("_T")
+    _T1 = TypeVar("_T1")
+    _T2 = TypeVar("_T2")
 
 
 logger = logging.getLogger(__name__)
@@ -76,60 +77,59 @@ def transform_v1_url(url: yarl.URL) -> yarl.URL:
 
 
 @overload
-async def throttled_gather(
-    coros: Iterable[Awaitable[_T]],
-    batch_size: int = 10,
+async def async_map(
+    coro_factory: Callable[[_T1], Awaitable[_T2]],
+    values: Iterable[_T1],
     *,
-    return_exceptions: Literal[False],
-) -> list[_T]: ...
+    return_exceptions: Literal[True],
+    task_limit: int = 10,
+) -> list[_T2 | Exception]: ...
 
 
 @overload
-async def throttled_gather(
-    coros: Iterable[Awaitable[_T]],
-    batch_size: int = 10,
+async def async_map(
+    coro_factory: Callable[[_T1], Awaitable[_T2]],
+    values: Iterable[_T1],
     *,
-    return_exceptions: bool = True,
-) -> list[_T | Exception]: ...
+    return_exceptions: Literal[False] = False,
+    task_limit: int = 10,
+) -> list[_T2]: ...
 
 
-async def throttled_gather(
-    coros: Iterable[Awaitable[_T]],
-    batch_size: int = 10,
+async def async_map(
+    coro_factory: Callable[[_T1], Awaitable[_T2]],
+    values: Iterable[_T1],
     *,
-    return_exceptions: bool = True,
-) -> Sequence[_T | Exception]:
-    """Like `asyncio.gather`, but creates tasks lazily to minimize event loop overhead.
+    return_exceptions: bool = False,
+    task_limit: int = 10,
+) -> Sequence[_T2 | Exception]:
+    """Creates tasks lazily to minimize event loop overhead.
 
-    This function ensures there are never more than `batch_size` tasks created at any given time.
-
-    If `return_exceptions` is `False`, any exceptions other than `asyncio.CancelledError` raised within
-    a task will cancel all remaining tasks and wait for them to exit.
-    The exceptions are then combined and raised as an `ExceptionGroup`.
+    This function ensures there are never more than `task_limit` tasks are created at any given time.
     """
-    semaphore = asyncio.BoundedSemaphore(batch_size)
-    tasks: list[asyncio.Task[_T | Exception]] = []
-
-    abort = False
+    semaphore = asyncio.BoundedSemaphore(task_limit)
+    tasks: list[asyncio.Task[_T2 | Exception]] = []
+    abort = asyncio.Event()
 
-    async def worker(coro: Awaitable[_T]) -> _T | Exception:
+    async def worker(coro: Awaitable[_T2]) -> _T2 | Exception:
         try:
             return await coro
         except Exception as e:
             if return_exceptions:
                 return e
-            nonlocal abort
-            abort = True
+            abort.set()
             raise
 
         finally:
             semaphore.release()
 
     async with asyncio.TaskGroup() as tg:
-        for coro in coros:
-            if abort:
+        for value in values:
+            if abort.is_set():
                 break
+
             await semaphore.acquire()
+            coro = coro_factory(value)
             tasks.append(tg.create_task(worker(coro)))
 
     return [t.result() for t in tasks]

async_mega_py-2.0.1.dev0/src/mega/download.py

@@ -1,87 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-import contextlib
-import errno
-import logging
-import shutil
-import tempfile
-from pathlib import Path
-from typing import IO, TYPE_CHECKING
-
-from mega import progress
-from mega.chunker import MegaChunker
-from mega.crypto import get_chunks
-
-if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator
-
-    import aiohttp
-
-
-logger = logging.getLogger(__name__)
-
-
-async def encrypted_stream(
-    stream: aiohttp.StreamReader,
-    output_path: Path,
-    file_size: int,
-    iv: tuple[int, int],
-    meta_mac: tuple[int, int],
-    key: tuple[int, int, int, int],
-) -> Path:
-    if await asyncio.to_thread(output_path.exists):
-        raise FileExistsError(errno.EEXIST, output_path)
-
-    chunker = MegaChunker(iv, key, meta_mac)
-    progress_hook = progress.current_hook.get()
-    async with _new_temp_download(output_path) as output:
-        for _, chunk_size in get_chunks(file_size):
-            encrypted_chunk = await stream.readexactly(chunk_size)
-            chunk = chunker.read(encrypted_chunk)
-            output.write(chunk)
-            progress_hook(len(chunk))
-
-        chunker.check_integrity()
-
-    return output_path
-
-
-async def stream(stream: aiohttp.StreamReader, output_path: Path) -> Path:
-    if await asyncio.to_thread(output_path.exists):
-        raise FileExistsError(errno.EEXIST, output_path)
-
-    chunk_size = 1024 * 1024 * 5  # 5MB
-    progress_hook = progress.current_hook.get()
-    async with _new_temp_download(output_path) as output:
-        async for chunk in stream.iter_chunked(chunk_size):
-            output.write(chunk)
-            progress_hook(len(chunk))
-
-    return output_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)
-    logger.info(f'Created temp file "{temp_file.name!s}" for "{output_path!s}"')
-    try:
-        yield temp_file
-
-        def move():
-            temp_file.close()
-            output_path.parent.mkdir(parents=True, exist_ok=True)
-            shutil.move(temp_file.name, output_path)
-            logger.info(f'Moved temp file "{temp_file.name!s}" to "{output_path!s}"')
-
-        await asyncio.to_thread(move)
-
-    finally:
-
-        def delete():
-            if not temp_file.closed:
-                temp_file.close()
-            Path(temp_file.name).unlink(missing_ok=True)
-
-        await asyncio.to_thread(delete)