refactor: update API class

pyproject.toml

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

src/mega/__init__.py

@@ -1,4 +1,8 @@
 import importlib.metadata
+from contextvars import ContextVar
 
 _package_name_ = "async-mega-py"
 __version__ = importlib.metadata.version(_package_name_)
+
+LOG_FILE_PROGRESS: ContextVar[bool] = ContextVar("LOG_FILE_PROGRESS", default=True)
+LOG_HTTP_TRAFFIC: ContextVar[bool] = ContextVar("LOG_HTTP_TRAFFIC", default=False)

src/mega/__main__.py

@@ -14,8 +14,7 @@
 import yarl
 from cyclopts import App, Parameter
 
-from mega import __version__, env
-from mega.api import LOG_HTTP_TRAFFIC
+from mega import LOG_HTTP_TRAFFIC, __version__, env
 from mega.client import MegaNzClient
 from mega.transfer_it import TransferItClient
 from mega.utils import Site, setup_logger

src/mega/api.py

@@ -7,15 +7,14 @@
 import logging
 import uuid
 from collections.abc import Mapping, Sequence
-from contextvars import ContextVar
 from functools import wraps
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, ParamSpec, Self, TypeVar
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, ParamSpec, Self, TypeVar
 
 import aiohttp
 import yarl
 from aiolimiter import AsyncLimiter
 
-from mega import __version__, _package_name_
+from mega import LOG_HTTP_TRAFFIC, __version__, _package_name_
 from mega.crypto import generate_hashcash
 from mega.errors import RequestError, RetryRequestError
 from mega.utils import random_id, random_u32int
@@ -27,8 +26,6 @@
     _R = TypeVar("_R")
 
 
-LOG_HTTP_TRAFFIC: ContextVar[bool] = ContextVar("LOG_HTTP_TRAFFIC", default=False)
-
 logger = logging.getLogger(__name__)
 
 
@@ -66,27 +63,23 @@ async def inner_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
     return wrapper
 
 
-@dataclasses.dataclass(slots=True, weakref_slot=True, init=False)
+@dataclasses.dataclass(slots=True, weakref_slot=True)
 class MegaAPI:
-    session_id: str | None
-    _request_id: int
-    _client_id: str
-
-    __session: aiohttp.ClientSession | None
-    _auto_close_session: bool
-    _rate_limiter: AsyncLimiter
-
-    _entrypoint: ClassVar[yarl.URL] = yarl.URL("https://g.api.mega.co.nz/cs")
-    user_agent: str
-
-    def __init__(self, session: aiohttp.ClientSession | None = None, user_agent: str | None = None) -> None:
-        self.session_id = None
-        self._request_id = random_u32int()
-        self._client_id = random_id(10)
-        self.__session = session
-        self._auto_close_session = session is None
-        self._rate_limiter = AsyncLimiter(100, 60)
-        self.user_agent = user_agent or f"{_package_name_}/{__version__}"
+    _session: aiohttp.ClientSession | None = None
+
+    user_agent: str = f"{_package_name_}/{__version__}"
+
+    session_id: str | None = dataclasses.field(init=False, default=None)
+    _request_id: int = dataclasses.field(init=False, default_factory=random_u32int)
+    _client_id: str = dataclasses.field(init=False, default_factory=lambda: random_id(10))
+
+    _auto_close_session: bool = dataclasses.field(init=False)
+    _rate_limiter: AsyncLimiter = dataclasses.field(init=False, default_factory=lambda: AsyncLimiter(100, 60))
+
+    _entrypoint: ClassVar[yarl.URL] = dataclasses.field(init=False, default=yarl.URL("https://g.api.mega.co.nz/cs"))
+
+    def __post_init__(self) -> None:
+        self._auto_close_session = self._session is None
 
     @property
     def entrypoint(self) -> yarl.URL:
@@ -96,22 +89,12 @@ def entrypoint(self) -> yarl.URL:
     def client_id(self) -> str:
         return self._client_id
 
-    @property
-    def request_id(self) -> int:
-        return self._request_id
-
-    @property
-    def _session(self) -> aiohttp.ClientSession:
-        if self.__session is None:
-            self.__session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(sock_connect=160, sock_read=60))
-        return self.__session
-
     def __repr__(self) -> str:
         return f"<{type(self).__name__}>(session_id={self.session_id!r}, client_id={self.client_id!r}, auto_close_session={self._auto_close_session!r})"
 
     async def aclose(self) -> None:
-        if self._auto_close_session and self.__session:
-            await self.__session.close()
+        if self._auto_close_session and self._session:
+            await self._session.close()
 
     async def __enter__(self) -> Self:
         return self
@@ -175,7 +158,7 @@ async def __request(
         headers: Mapping[str, str] | None = None,
         **kwargs: Any,
     ) -> AsyncGenerator[aiohttp.ClientResponse]:
-        kwargs["headers"] = {"User-Agent": self.user_agent} | (headers or {})
+        kwargs["headers"] = {"User-Agent": self.user_agent, **(headers or {})}
         request_id = str(uuid.uuid4())
         if LOG_HTTP_TRAFFIC.get():
             logger.debug(
@@ -187,6 +170,9 @@ async def __request(
             )
 
         resp = None
+        if self._session is None:
+            self._session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(sock_connect=160, sock_read=60))
+
         try:
             async with self._rate_limiter, self._session.request(method, url, **kwargs) as resp:
                 yield resp
@@ -244,19 +230,19 @@ def __str__(self) -> str:
         return str(self.__json__())
 
 
-class APIContextManager:
-    __slots__ = ("_api",)
+_API_T = TypeVar("_API_T", bound=MegaAPI, covariant=True)
 
-    def __init__(self, session: aiohttp.ClientSession | None = None, *, user_agent: str | None = None) -> None:
-        self._api: MegaAPI = MegaAPI(session, user_agent)
+
+class APIContextManager(Generic[_API_T]):
+    __slots__ = ("_api",)
 
     async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(self, *_) -> None:
-        await self.close()
+        await self.aclose()
 
     async def aclose(self) -> None:
-        await self._api.close()
+        await self._api.aclose()
 
     close = aclose

src/mega/client.py

@@ -7,7 +7,7 @@
 from typing import TYPE_CHECKING
 
 from mega import progress
-from mega.api import APIContextManager
+from mega.api import APIContextManager, MegaAPI
 from mega.core import MegaCore
 from mega.crypto import a32_to_base64, b64_to_a32, b64_url_encode, encrypt_attr, encrypt_key
 from mega.data_structures import (
@@ -38,13 +38,16 @@
 _DOMAIN = Site.MEGA.value
 
 
-class MegaNzClient(APIContextManager):
+class MegaNzClient(APIContextManager[MegaAPI]):
     """Interface with all the public methods of the API"""
 
     __slots__ = ("_core",)
 
     def __init__(self, session: aiohttp.ClientSession | None = None, *, user_agent: str | None = None) -> None:
-        super().__init__(session, user_agent=user_agent)
+        self._api = MegaAPI(session)
+        if user_agent:
+            self._api.user_agent = user_agent
+
         self._core: MegaCore = MegaCore(self._api)
         if hasattr(sys, "ps1"):
             setup_logger(logging.DEBUG)

src/mega/transfer_it.py

@@ -45,9 +45,13 @@ async def post(self, json: dict[str, Any] | list[dict[str, Any]], params: dict[s
         return await super().post(json, params)
 
 
-class TransferItClient(APIContextManager):
-    def __init__(self, session: aiohttp.ClientSession | None = None, *, user_agent: str | None = None) -> None:  # pyright: ignore[reportMissingSuperCall]
-        self._api = TransferItAPI(session, user_agent=user_agent)
+class TransferItClient(APIContextManager[TransferItAPI]):
+    __slots__ = ()
+
+    def __init__(self, session: aiohttp.ClientSession | None = None, *, user_agent: str | None = None) -> None:
+        self._api = TransferItAPI(session)
+        if user_agent:
+            self._api.user_agent = user_agent
 
     @property
     def progress_bar(self) -> _GeneratorContextManager[None, None, None]:

src/mega/utils.py

@@ -14,6 +14,7 @@
 import aiohttp
 import yarl
 
+from mega import LOG_FILE_PROGRESS
 from mega.errors import ValidationError
 
 if TYPE_CHECKING:
@@ -63,7 +64,7 @@ def setup_logger(level: int = logging.INFO) -> None:
 
 
 def progress_logger(output_path: Path, file_size: int, *, download: bool) -> Callable[[float], None]:
-    if not logger.isEnabledFor(10):
+    if not (LOG_FILE_PROGRESS.get() and logger.isEnabledFor(logging.DEBUG)):
         return lambda _: None
 
     from mega.data_structures import ByteSize