|
|
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, |
|
|
}, |
|
|
) |
|
|
|