PyNyaaTa/pynyaata2/types.py

from abc import ABC, abstractmethod
from asyncio import get_event_loop
from datetime import datetime
from enum import Enum
from functools import partial, wraps
from typing import List, Optional

from pydantic import BaseModel, ByteSize, HttpUrl


class Color(str, Enum):
    WHITE = "white"
    BLACK = "black"
    LIGHT = "light"
    DARK = "dark"
    PRIMARY = "primary"
    LINK = "link"
    INFO = "info"
    SUCCESS = "success"
    WARNING = "warning"
    DANGER = "danger"
    DEFAULT = "default"


class RemoteFile(BaseModel):
    bridge: str
    id: int
    category: str
    color: Optional[Color]
    name: str
    link: HttpUrl
    comment: int = 0
    comment_url: Optional[HttpUrl]
    magnet: Optional[str]
    torrent: Optional[HttpUrl]
    size: Optional[ByteSize]
    date: Optional[datetime]
    seeds: Optional[int]
    leechs: Optional[int]
    downloads: Optional[int]
    nb_pages: int = 1


class Bridge(BaseModel, ABC):
    color: Color
    title: str
    favicon: HttpUrl
    base_url: HttpUrl

    @abstractmethod
    def search_url(self, query: str = "", page: int = 1) -> HttpUrl:
        pass

    @abstractmethod
    async def search(self, query: str = "", page: int = 1) -> List[RemoteFile]:
        pass


class Cache(ABC):
    @abstractmethod
    def get(self, key: str) -> Optional[List[RemoteFile]]:
        pass

    @abstractmethod
    def set(self, key: str, data: List[RemoteFile]):
        pass


def async_wrap(func):
    @wraps(func)
    async def run(*args, loop=None, executor=None, **kwargs):
        if loop is None:
            loop = get_event_loop()

        pfunc = partial(func, *args, **kwargs)
        return await loop.run_in_executor(executor, pfunc)

    return run