zhenxun_bot/models/pixiv.py

from typing import Optional, List
from services.db_context import db
import asyncio


class Pixiv(db.Model):
    __tablename__ = "pixiv"

    id = db.Column(db.Integer(), primary_key=True)
    pid = db.Column(db.BigInteger(), nullable=False)
    title = db.Column(db.String(), nullable=False)
    width = db.Column(db.Integer(), nullable=False)
    height = db.Column(db.Integer(), nullable=False)
    view = db.Column(db.Integer(), nullable=False)
    bookmarks = db.Column(db.Integer(), nullable=False)
    img_url = db.Column(db.String(), nullable=False)
    img_p = db.Column(db.String(), nullable=False)
    uid = db.Column(db.BigInteger(), nullable=False)
    author = db.Column(db.String(), nullable=False)
    is_r18 = db.Column(db.Boolean(), nullable=False)
    tags = db.Column(db.String(), nullable=False)

    _idx1 = db.Index("pixiv_idx1", "pid", "img_url", unique=True)

    @classmethod
    async def add_image_data(
            cls,
            pid: int,
            title: str,
            width: int,
            height: int,
            view: int,
            bookmarks: int,
            img_url: str,
            img_p: str,
            uid: int,
            author: str,
            tags: str,
    ):
        """
        说明：
            添加图片信息
        参数：
            :param pid: pid
            :param title: 标题
            :param width: 宽度
            :param height: 长度
            :param view: 被查看次数
            :param bookmarks: 收藏数
            :param img_url: url链接
            :param img_p: 张数
            :param uid: 作者uid
            :param author: 作者名称
            :param tags: 相关tag
        """
        if not await cls.check_exists(pid, img_p):
            await cls.create(
                pid=pid,
                title=title,
                width=width,
                height=height,
                view=view,
                bookmarks=bookmarks,
                img_url=img_url,
                img_p=img_p,
                uid=uid,
                author=author,
                is_r18=True if 'R-18' in tags else False,
                tags=tags,
            )
            return True
        return False

    @classmethod
    async def remove_image_data(cls, pid: int, img_p: str) -> bool:
        """
        说明：
            删除图片数据
        参数：
            :param pid: 图片pid
            :param img_p: 图片pid的张数，如：p0，p1
        """
        try:
            if img_p:
                await cls.delete.where(
                    (cls.pid == pid) & (cls.img_p == img_p)
                ).gino.status()
            else:
                await cls.delete.where(cls.pid == pid).gino.status()
            return True
        except Exception:
            return False

    @classmethod
    async def get_all_pid(cls) -> List[int]:
        """
        说明：
            获取所有PID
        """
        pid = []
        query = await cls.query.gino.all()
        for image in query:
            if image.pid not in pid:
                pid.append(image.pid)
        return pid

    # 0：非r18  1：r18  2：混合
    @classmethod
    async def query_images(
        cls,
        keyword: Optional[List[str]] = None,
        uid: Optional[int] = None,
        pid: Optional[int] = None,
        r18: int = 0,
    ) -> List[Optional["Pixiv"]]:
        """
        说明：
            查找符合条件的图片
        参数：
            :param keyword: 关键词
            :param uid: 画师uid
            :param pid: 图片pid
            :param r18: 是否r18，0：非r18  1：r18  2：混合
        """
        if r18 == 0:
            query = await cls.query.where(cls.is_r18 == False).gino.all()
        elif r18 == 1:
            query = await cls.query.where(cls.is_r18 == True).gino.all()
        else:
            query = await cls.query.gino.all()
        if keyword:
            query = [x for x in query if set(x.tags.split(',')) > set(keyword)]
        elif uid:
            query = [x for x in query if x.uid == uid]
        elif pid:
            query = [x for x in query if x.pid == pid]
        return query

    @classmethod
    async def check_exists(cls, pid: int, img_p: str) -> bool:
        """
        说明：
            检测pid是否已存在
        参数：
            :param pid: 图片PID
            :param img_p: 张数
        """
        query = await cls.query.where(
            (cls.pid == pid) & (cls.img_p == img_p)
        ).gino.all()
        return bool(query)

    @classmethod
    async def get_keyword_num(cls, keyword: List[str]) -> "int, int":
        """
        说明：
            获取相关关键词(keyword, tag)在图库中的数量
        参数：
            :param keyword: 关键词/Tag
        """
        query = await cls.query.gino.all()
        i = int(len(query) / 200)
        mod = len(query) % 200
        tasks = []
        start = 0
        end = 200
        count = 0
        r18_count = 0
        for _ in range(i):
            tasks.append(asyncio.ensure_future(split_query_list(query[start: end], keyword)))
            start += 200
            end += 200
        if mod:
            tasks.append(asyncio.ensure_future(split_query_list(query[end:], keyword)))
        result = await asyncio.gather(*tasks)
        for x, j in result:
            count += x
            r18_count += j
        # query = [x for x in query if set(x.tags.split(',')) > set(keyword)]
        # r18_count = len([x for x in query if x.is_r18])
        return count, r18_count


async def split_query_list(query: List['Pixiv'], keyword: List[str]) -> 'int, int':
    return await asyncio.get_event_loop().run_in_executor(None, _split_query_list, query, keyword)


def _split_query_list(query: List['Pixiv'], keyword: List[str]) -> 'int, int':
    query = [x for x in query if set(x.tags.split(',')) > set(keyword)]
    r18_count = len([x for x in query if x.is_r18])
    return len(query) - r18_count, r18_count