beatapi.py

"""Small dependency-free BeatAPI reference client.

This module is example code, not an officially versioned SDK.
"""

from __future__ import annotations

import json
import os
import random as random_module
import time
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, Callable


TERMINAL_STATUSES = {"succeeded", "failed"}


class BeatAPIError(Exception):
    def __init__(
        self,
        message: str,
        *,
        status: int | None = None,
        code: str | None = None,
        request_id: str | None = None,
        details: Any = None,
    ) -> None:
        super().__init__(message)
        self.status = status
        self.code = code
        self.request_id = request_id
        self.details = details


class BeatAPIClient:
    def __init__(
        self,
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        transport: Callable[[urllib.request.Request], Any] = urllib.request.urlopen,
        sleep: Callable[[float], None] = time.sleep,
        random: Callable[[], float] = random_module.random,
    ) -> None:
        self.api_key = api_key if api_key is not None else os.getenv("BEATAPI_API_KEY")
        if not self.api_key:
            raise ValueError(
                "BEATAPI_API_KEY is required. Create a key in the BeatAPI dashboard."
            )

        self.base_url = (
            base_url or os.getenv("BEATAPI_BASE_URL") or "https://api.beatapi.io"
        ).rstrip("/")
        self.transport = transport
        self.sleep = sleep
        self.random = random

    def request(
        self,
        path: str,
        *,
        method: str = "GET",
        body: dict[str, Any] | None = None,
    ) -> Any:
        data = json.dumps(body).encode() if body is not None else None
        headers = {
            "Accept": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }
        if data is not None:
            headers["Content-Type"] = "application/json"

        request = urllib.request.Request(
            f"{self.base_url}{path}",
            data=data,
            headers=headers,
            method=method,
        )

        try:
            response = self.transport(request)
            with response:
                status = response.status
                payload = json.loads(response.read().decode() or "{}")
        except urllib.error.HTTPError as error:
            status = error.code
            payload = json.loads(error.read().decode() or "{}")

        if status < 200 or status >= 300:
            public_error = payload.get("error", {})
            raise BeatAPIError(
                public_error.get(
                    "message", f"BeatAPI request failed with HTTP {status}."
                ),
                status=status,
                code=public_error.get("code"),
                request_id=public_error.get("request_id"),
                details=public_error.get("details"),
            )

        return payload.get("data")

    def create_music_video_task(self, input_data: dict[str, Any]) -> dict[str, Any]:
        return self.request(
            "/v1/music-video/tasks",
            method="POST",
            body=input_data,
        )

    def create_ecommerce_video_task(
        self, input_data: dict[str, Any]
    ) -> dict[str, Any]:
        return self.request(
            "/v1/ecommerce-video/tasks",
            method="POST",
            body=input_data,
        )

    def get_task(self, task_id: str) -> dict[str, Any]:
        return self.request(f"/v1/tasks/{urllib.parse.quote(task_id, safe='')}")

    def get_usage(self) -> dict[str, Any]:
        return self.request("/v1/usage")

    def wait_for_task(
        self,
        task_id: str,
        *,
        interval_seconds: float = 5,
        max_attempts: int = 120,
        on_update: Callable[[dict[str, Any], int], None] | None = None,
    ) -> dict[str, Any]:
        for attempt in range(1, max_attempts + 1):
            task = self.get_task(task_id)
            if on_update:
                on_update(task, attempt)

            if task["status"] in TERMINAL_STATUSES:
                return task

            if attempt < max_attempts:
                jitter = interval_seconds * 0.2 * self.random()
                self.sleep(interval_seconds + jitter)

        raise TimeoutError(
            f"Task {task_id} did not finish after {max_attempts} attempts."
        )