six2one 0.1.0__py3-none-any.whl

six2one/__init__.py

@@ -0,0 +1,3 @@
+"""Manifest-backed e621/e926 dataset fetching."""
+
+__version__ = "0.1.0"

six2one/api.py

@@ -0,0 +1,132 @@
+"""Small async client for the e621/e926 endpoints used by fetch."""
+
+from pathlib import PurePosixPath
+from typing import Any, Final
+from urllib.parse import urlparse
+
+from tqdm import tqdm
+
+from .auth import load_login, request_headers
+from .models import Site
+from .network import RequestAdapter
+
+
+MAX_POSTS_PER_REQUEST: Final = 320
+MAX_TAGS_PER_REQUEST: Final = 320
+DOWNLOAD_CHUNK_SIZE_BYTES: Final = 64 * 1024
+
+
+class E621API:
+    """Async API client for e621-compatible GET operations."""
+
+    def __init__(self, site: Site = Site.E621) -> None:
+        self.site = site
+        self.base_url = site.base_url
+        self._adapter: RequestAdapter | None = None
+
+    async def __aenter__(self) -> "E621API":
+        self._adapter = RequestAdapter(headers=request_headers(load_login()))
+        await self._adapter.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        if self._adapter is not None:
+            await self._adapter.__aexit__(exc_type, exc_val, exc_tb)
+            self._adapter = None
+
+    async def get_posts(
+        self,
+        tags: str,
+        limit: int,
+        page: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """Search posts.
+
+        Raises:
+            ValueError: If limit is outside the e621 per-request bounds.
+            RuntimeError: If called outside the async context manager.
+            aiohttp.ClientResponseError: If the site returns an error response.
+        """
+        _validate_limit(limit, MAX_POSTS_PER_REQUEST, "post")
+        params: dict[str, str] = {
+            "tags": tags,
+            "limit": str(limit),
+        }
+        if page is not None:
+            params["page"] = page
+        data = await self._get("/posts.json", params)
+        if "posts" not in data:
+            raise ValueError("Post search response did not include 'posts'")
+        posts = data["posts"]
+        if not isinstance(posts, list):
+            raise ValueError("Post search response 'posts' value was not a list")
+        return posts
+
+    async def get_tags(self, name_matches: str, limit: int = 1) -> list[dict[str, Any]]:
+        """List tags matching a name expression.
+
+        Raises:
+            ValueError: If limit is outside the e621 per-request bounds.
+            RuntimeError: If called outside the async context manager.
+            aiohttp.ClientResponseError: If the site returns an error response.
+        """
+        _validate_limit(limit, MAX_TAGS_PER_REQUEST, "tag")
+        params = {
+            "search[name_matches]": name_matches,
+            "limit": str(limit),
+        }
+        data = await self._get("/tags.json", params)
+        if not isinstance(data, list):
+            raise ValueError("Tag listing response was not a list")
+        return data
+
+    async def download_url(self, url: str) -> bytes:
+        """Download a URL with a byte progress bar and return its response body.
+
+        Raises:
+            RuntimeError: If called outside the async context manager.
+            aiohttp.ClientResponseError: If the site returns an error response.
+        """
+        adapter = self._require_adapter()
+        request_context = await adapter.request("GET", url)
+        async with request_context as response:
+            response.raise_for_status()
+            chunks: list[bytes] = []
+            with tqdm(
+                total=response.content_length,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                desc=_download_description(url),
+            ) as progress_bar:
+                async for chunk in response.content.iter_chunked(DOWNLOAD_CHUNK_SIZE_BYTES):
+                    chunks.append(chunk)
+                    progress_bar.update(len(chunk))
+            return b"".join(chunks)
+
+    async def _get(self, endpoint: str, params: dict[str, str]) -> Any:
+        adapter = self._require_adapter()
+        url = f"{self.base_url}{endpoint}"
+        request_context = await adapter.request("GET", url, params=params)
+        async with request_context as response:
+            response.raise_for_status()
+            return await response.json()
+
+    def _require_adapter(self) -> RequestAdapter:
+        if self._adapter is None:
+            raise RuntimeError("E621API must be used as an async context manager")
+        return self._adapter
+
+
+def _validate_limit(limit: int, maximum: int, label: str) -> None:
+    if limit < 0:
+        raise ValueError(f"{label} request limit must be at least zero")
+    if limit > maximum:
+        raise ValueError(f"{label} request limit cannot exceed {maximum}")
+
+
+def _download_description(url: str) -> str:
+    path_name = PurePosixPath(urlparse(url).path).name
+    if path_name:
+        return path_name
+    return "download"

six2one/auth.py

@@ -0,0 +1,131 @@
+"""Project-local e621 login storage and request headers."""
+
+import base64
+from dataclasses import dataclass
+import json
+import os
+from pathlib import Path
+from typing import Final
+
+from .errors import UsageError
+from .models import TOOL_NAME, TOOL_VERSION
+
+
+LOGIN_FILENAME: Final = ".six2one-login.json"
+PROJECT_MARKER_FILENAME: Final = "pyproject.toml"
+
+
+@dataclass(frozen=True)
+class LoginCredentials:
+    """Validated login credentials for e621-compatible Basic auth."""
+
+    username: str
+    api_key: str
+
+
+def find_project_root(start_path: Path | None = None) -> Path:
+    """Find the nearest project root from a start path.
+
+    Raises:
+        UsageError: If no project marker is found.
+    """
+    current_path = Path.cwd() if start_path is None else start_path
+    resolved_path = current_path.resolve()
+    search_path = resolved_path if resolved_path.is_dir() else resolved_path.parent
+    for candidate in (search_path, *search_path.parents):
+        if (candidate / PROJECT_MARKER_FILENAME).is_file():
+            return candidate
+    raise UsageError(
+        f"Could not find project root from {resolved_path}; missing {PROJECT_MARKER_FILENAME}"
+    )
+
+
+def login_path(project_root: Path | None = None) -> Path:
+    """Return the login file path for the project root."""
+    root = find_project_root() if project_root is None else project_root
+    return root / LOGIN_FILENAME
+
+
+def save_login(username: str, api_key: str, project_root: Path | None = None) -> Path:
+    """Save login credentials to the project root.
+
+    Raises:
+        UsageError: If username or api key is empty.
+    """
+    credentials = _credentials_from_values(username, api_key)
+    path = login_path(project_root)
+    data = {
+        "username": credentials.username,
+        "api_key": credentials.api_key,
+    }
+    with path.open("w", encoding="utf-8") as file:
+        json.dump(data, file, indent=2, sort_keys=True)
+        file.write("\n")
+    os.chmod(path, 0o600)
+    return path
+
+
+def delete_login(project_root: Path | None = None) -> Path:
+    """Delete the project login file.
+
+    Raises:
+        UsageError: If the login file does not exist.
+    """
+    path = login_path(project_root)
+    if not path.exists():
+        raise UsageError(f"No login file exists at {path}")
+    path.unlink()
+    return path
+
+
+def load_login(project_root: Path | None = None) -> LoginCredentials | None:
+    """Load project login credentials if present.
+
+    Raises:
+        UsageError: If the login file is malformed.
+    """
+    path = login_path(project_root)
+    if not path.exists():
+        return None
+    with path.open("r", encoding="utf-8") as file:
+        try:
+            data = json.load(file)
+        except json.JSONDecodeError as error:
+            raise UsageError(f"Login file is not valid JSON: {path}") from error
+    if not isinstance(data, dict):
+        raise UsageError(f"Login file root must be a JSON object: {path}")
+    if "username" not in data:
+        raise UsageError(f"Login file is missing required key: username")
+    if "api_key" not in data:
+        raise UsageError(f"Login file is missing required key: api_key")
+    return _credentials_from_values(data["username"], data["api_key"])
+
+
+def request_headers(credentials: LoginCredentials | None) -> dict[str, str]:
+    """Build HTTP headers for API requests."""
+    if credentials is None:
+        return {
+            "User-Agent": f"{TOOL_NAME}/{TOOL_VERSION}",
+        }
+    token = base64.b64encode(f"{credentials.username}:{credentials.api_key}".encode("utf-8"))
+    return {
+        "Authorization": f"Basic {token.decode('ascii')}",
+        "User-Agent": f"{TOOL_NAME}/{TOOL_VERSION} (by {credentials.username} on e621)",
+    }
+
+
+def _credentials_from_values(username: object, api_key: object) -> LoginCredentials:
+    if not isinstance(username, str):
+        raise UsageError("username must be a string")
+    if not isinstance(api_key, str):
+        raise UsageError("api_key must be a string")
+    normalized_username = username.strip()
+    normalized_api_key = api_key.strip()
+    if not normalized_username:
+        raise UsageError("username cannot be empty")
+    if not normalized_api_key:
+        raise UsageError("api_key cannot be empty")
+    return LoginCredentials(
+        username=normalized_username,
+        api_key=normalized_api_key,
+    )