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]


@dataclass_json
@dataclass
class E621PostVersion:
    id: int
    post_id: int
    tags: str
    updater_id: int
    updated_at: str
    rating: Rating
    parent_id: Optional[int]
    source: Optional[str]
    description: str
    reason: Optional[str]
    locked_tags: Optional[str]
    added_tags: List[str]
    removed_tags: List[str]
    added_locked_tags: List[str]
    removed_locked_tags: List[str]
    rating_changed: bool
    parent_changed: bool
    source_changed: bool
    description_changed: bool
    version: int
    obsolete_added_tags: str
    obsolete_removed_tags: str
    unchanged_tags: str
    updater_name: str


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()
        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]

    async def get_post_versions(self, start_id=0, page=1, limit=320) -> List[E621PostVersion]:
        r = (await self.client.get('/post_versions.json', params={
            'search[start_id]': start_id,
            'limit': limit,
            'page': page,
        })).json()
        if 'success' in r:
            return []
        return [E621PostVersion.from_dict(p) for p in r]