hibiapi/api/sauce/api.py

import random
from enum import IntEnum
from io import BytesIO
from typing import Any, Optional, overload

from httpx import HTTPError

from hibiapi.api.sauce.constants import SauceConstants
from hibiapi.utils.decorators import enum_auto_doc
from hibiapi.utils.exceptions import ClientSideException
from hibiapi.utils.net import catch_network_error
from hibiapi.utils.routing import BaseEndpoint, BaseHostUrl


class UnavailableSourceException(ClientSideException):
    code = 422
    detail = "given image is not avaliable to fetch"


class ImageSourceOversizedException(UnavailableSourceException):
    code = 413
    detail = (
        "given image size is rather than maximum limit "
        f"{SauceConstants.IMAGE_MAXIMUM_SIZE} bytes"
    )


class HostUrl(BaseHostUrl):
    allowed_hosts = SauceConstants.IMAGE_ALLOWED_HOST


class UploadFileIO(BytesIO):
    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v: Any) -> BytesIO:
        if not isinstance(v, BytesIO):
            raise ValueError(f"Expected UploadFile, received: {type(v)}")
        return v


@enum_auto_doc
class DeduplicateType(IntEnum):
    DISABLED = 0
    """no result deduplicating"""
    IDENTIFIER = 1
    """consolidate search results and deduplicate by item identifier"""
    ALL = 2
    """all implemented deduplicate methods such as by series name"""


class SauceEndpoint(BaseEndpoint, cache_endpoints=False):
    base = "https://saucenao.com"

    async def fetch(self, host: HostUrl) -> UploadFileIO:
        try:
            response = await self.client.get(
                url=host,
                headers=SauceConstants.IMAGE_HEADERS,
                timeout=SauceConstants.IMAGE_TIMEOUT,
            )
            response.raise_for_status()
            if len(response.content) > SauceConstants.IMAGE_MAXIMUM_SIZE:
                raise ImageSourceOversizedException
            return UploadFileIO(response.content)
        except HTTPError as e:
            raise UnavailableSourceException(detail=str(e)) from e

    @catch_network_error
    async def request(
        self, *, file: UploadFileIO, params: dict[str, Any]
    ) -> dict[str, Any]:
        response = await self.client.post(
            url=self._join(
                self.base,
                "search.php",
                params={
                    **params,
                    "api_key": random.choice(SauceConstants.API_KEY),
                    "output_type": 2,
                },
            ),
            files={"file": file},
        )
        if response.status_code >= 500:
            response.raise_for_status()
        return response.json()

    @overload
    async def search(
        self,
        *,
        url: HostUrl,
        size: int = 30,
        deduplicate: DeduplicateType = DeduplicateType.ALL,
        database: Optional[int] = None,
        enabled_mask: Optional[int] = None,
        disabled_mask: Optional[int] = None,
    ) -> dict[str, Any]:
        ...

    @overload
    async def search(
        self,
        *,
        file: UploadFileIO,
        size: int = 30,
        deduplicate: DeduplicateType = DeduplicateType.ALL,
        database: Optional[int] = None,
        enabled_mask: Optional[int] = None,
        disabled_mask: Optional[int] = None,
    ) -> dict[str, Any]:
        ...

    async def search(
        self,
        *,
        url: Optional[HostUrl] = None,
        file: Optional[UploadFileIO] = None,
        size: int = 30,
        deduplicate: DeduplicateType = DeduplicateType.ALL,
        database: Optional[int] = None,
        enabled_mask: Optional[int] = None,
        disabled_mask: Optional[int] = None,
    ):
        if url is not None:
            file = await self.fetch(url)
        assert file is not None
        return await self.request(
            file=file,
            params={
                "dbmask": enabled_mask,
                "dbmaski": disabled_mask,
                "db": database,
                "numres": size,
                "dedupe": deduplicate,
            },
        )