e621_bot/e621.py

import json
import logging
from dataclasses import dataclass
from dataclasses_json import dataclass_json
from enum import Enum
from typing import List, Optional

import httpx


class Rating(Enum):
    E = "e"
    Q = "q"
    S = "s"


@dataclass_json
@dataclass
class E621PostFile:
    width: int
    height: int
    ext: str
    size: int
    md5: str
    url: Optional[str]


@dataclass_json
@dataclass
class E621PostFlags:
    pending: bool
    flagged: bool
    note_locked: bool
    status_locked: bool
    rating_locked: bool
    # comment_disabled: Optional[bool]
    deleted: bool


@dataclass_json
@dataclass
class E621PostPreview:
    width: int
    height: int
    url: Optional[str]


@dataclass_json
@dataclass
class E621PostRelationships:
    parent_id: Optional[int]
    has_children: bool
    has_active_children: bool
    children: List[int]


@dataclass_json
@dataclass
class E621PostOriginal:
    type: str
    height: int
    width: int
    urls: List[Optional[str]]


@dataclass_json
@dataclass
class E621PostAlternates:
    original: Optional[E621PostOriginal]


@dataclass_json
@dataclass
class E621PostSample:
    has: bool
    height: int
    width: int
    url: Optional[str]
    # alternates: E621PostAlternates


@dataclass_json
@dataclass
class E621PostScore:
    up: int
    down: int
    total: int


@dataclass_json
@dataclass
class E621PostTags:
    general: List[str]
    species: List[str]
    character: List[str]
    copyright: List[str]
    artist: List[str]
    invalid: List[str]
    lore: List[str]
    meta: List[str]

    def flatten(self):
        return self.general + self.species + self.character + self.copyright + self.artist + \
               self.invalid + self.lore + self.meta


@dataclass_json
@dataclass
class E621Post:
    id: int
    created_at: str
    updated_at: str
    file: E621PostFile
    preview: E621PostPreview
    sample: E621PostSample
    score: E621PostScore
    tags: E621PostTags
    locked_tags: List[str]
    change_seq: int
    flags: E621PostFlags
    rating: Rating
    fav_count: int
    sources: List[str]
    pools: List[int]
    relationships: E621PostRelationships
    approver_id: Optional[int]
    uploader_id: int
    description: str
    comment_count: int
    is_favorited: bool
    has_notes: bool
    duration: Optional[float]


class E621:
    def __init__(self):
        self.client = httpx.AsyncClient(headers={'user-agent': 'bot/1.0 (bakatrouble)'}, base_url='https://e621.net')

    async def get_posts(self, tags='', page=1, limit=320) -> List[E621Post]:
        r = (await self.client.get('/posts.json', params={'tags': tags, 'page': page, 'limit': limit})).json()
        logging.warning(json.dumps(r))
        return [E621Post.from_dict(p) for p in r['posts']]

    async def get_post(self, post_id: str) -> E621Post:
        return (await self.get_posts(f'id:{post_id}'))[0]