🚧

src/async_pixiv/client/api/_illust.py

@@ -1,12 +1,15 @@
 from pydantic import Field
 
 from async_pixiv.client.api._abc import APIBase
+from async_pixiv.const import APP_API_HOST
 from async_pixiv.model import Illust, PixivModel
 from async_pixiv.model.other.enums import SearchShort, SearchTarget
-from async_pixiv.model.other.result import PageResult
+from async_pixiv.model.other.result import PageResult, UgoiraMetadataResult
 
 __all__ = ("IllustAPI", "IllustPageResult", "IllustDetail")
 
+from async_pixiv.utils.context import set_pixiv_client
+
 
 class IllustDetail(PixivModel):
     illust: Illust
@@ -33,3 +36,12 @@ async def search(
         return await super().search(
             words, sort=sort, duration=duration, target=target, offset=offset, **kwargs
         )
+
+    # noinspection PyShadowingBuiltins
+    async def ugoira_metadata(self, id: int) -> UgoiraMetadataResult:
+        response = await self._pixiv_client.request_get(
+                APP_API_HOST / "v1/ugoira/metadata", params={"illust_id": id}
+        )
+        response.raise_for_result().raise_for_status()
+        with set_pixiv_client(self._pixiv_client):
+            return UgoiraMetadataResult.model_validate(response.json())

src/async_pixiv/client/api/_illust.pyi

@@ -3,7 +3,7 @@ from typing import Sequence
 from async_pixiv.client.api._abc import APIBase
 from async_pixiv.model import Illust, PixivModel
 from async_pixiv.model.other.enums import SearchShort, SearchTarget
-from async_pixiv.model.other.result import PageResult
+from async_pixiv.model.other.result import PageResult, UgoiraMetadataResult
 from async_pixiv.typedefs import DurationTypes, ShortTypes
 
 __all__ = ("IllustAPI", "IllustPageResult", "IllustDetail")
@@ -34,3 +34,6 @@ class IllustAPI(APIBase):
 
     async def recommended(self) -> IllustPageResult:
         pass
+
+    async def ugoira_metadata(self, id: int) -> UgoiraMetadataResult:
+        pass

src/async_pixiv/model/illust.py

@@ -1,14 +1,26 @@
+from asyncio import Event
 from functools import cached_property
-from typing import Any
+from io import BytesIO
+from pathlib import Path
+from typing import Any, Literal, overload
+from zipfile import ZipFile
 
+from aiofiles.tempfile import TemporaryDirectory
 from pydantic import Field
+from functools import cache
+from aiofiles import open as async_open
 
+from async_pixiv.error import ArtWorkTypeError
 from async_pixiv.model._base import PixivModel
+from async_pixiv.model.other.enums import Quality
 from async_pixiv.model.other.image import ImageUrl
+from async_pixiv.model.other.result import UgoiraMetadata
 from async_pixiv.model.other.tag import Tag
 from async_pixiv.model.user import User
 from async_pixiv.typedefs import Datetime, Enum, URL
+from async_pixiv.utils.ffmpeg import FFmpeg
 
+UGOIRA_RESULT_TYPE = Literal["zip", "gif", "mp4", "frame"]
 
 class IllustType(Enum):
     illust = "illust"
@@ -70,3 +82,127 @@ class Illust(PixivModel):
     @cached_property
     def link(self) -> URL:
         return URL(f"https://www.pixiv.net/artworks/{self.id}")
+
+    @cache
+    async def get_ugoira_metadata(self) -> UgoiraMetadata:
+        return (await self._pixiv_client.ILLUST.ugoira_metadata(self.id)).metadata
+
+    @overload
+    async def download_ugoira(self, quality: Quality = Quality.Original,*, result_type: Literal["zip"] = "zip") -> bytes | None:
+        """type of zip"""
+
+    @overload
+    async def download_ugoira(self, quality: Quality = Quality.Original,*, result_type: Literal["frame"] = "frame") -> list[bytes] | None:
+        """type of frame"""
+
+    @overload
+    async def download_ugoira(self, quality: Quality = Quality.Original,*, result_type: Literal["gif"] = "gif") -> bytes | None:
+        """type of GIF"""
+
+    @overload
+    async def download_ugoira(self, quality: Quality = Quality.Original,*, result_type: Literal["mp4"] = "mp4") -> bytes | None:
+        """type of mp4"""
+
+    async def download_ugoira(
+        self, quality: Quality = Quality.Original,*, result_type: UGOIRA_RESULT_TYPE = "zip"
+    ) -> bytes | list[bytes] | None:
+        if self.type != IllustType.ugoira:
+            raise ArtWorkTypeError(
+                "If you want to download a normal image, "
+                'please use this method: "download"'
+            )
+
+        metadata = await self.get_ugoira_metadata()
+
+        match quality:
+            case Quality.Large:
+                link = metadata.zip_url.large or metadata.zip_url.medium or metadata.zip_url.square
+            case Quality.Medium:
+                link = metadata.zip_url.medium or metadata.zip_url.square
+            case Quality.Square:
+                link = metadata.zip_url.square
+            case _:
+                link = metadata.zip_url.link
+        link = metadata.zip_url.link if link is None else link
+
+        data = await self._pixiv_client.download(link)
+        if data is None:
+            return None
+        if result_type == "zip":
+            return data
+
+        zip_file = ZipFile(BytesIO(data))
+
+        if result_type == "frame":
+            frames = []
+            for frame in metadata.frames:
+                with zip_file.open(frame.file) as f:
+                    frames.append(f.read())
+            return frames
+
+        async with TemporaryDirectory() as directory:
+            directory = Path(directory).resolve()
+            connect_config_file_path = directory / "list.txt"
+            async with async_open(connect_config_file_path, mode="w") as list_file:
+                for frame in metadata.frames:
+                    with zip_file.open(frame.file) as frame_zip_file:
+                        frame_file_path = directory / frame.file
+                        async with async_open(frame_file_path, mode="wb") as frame_file:
+                            await frame_file.write(frame_zip_file.read())
+                    await list_file.write(
+                        f"file {frame_file_path.resolve()}\n".replace("\\", "/")
+                    )
+                    await list_file.write(f"duration {frame.delay / 1000}\n")
+            del zip_file
+            event = Event()
+            if result_type == "mp4":
+                output_path = directory / "out.mp4"
+                # noinspection SpellCheckingInspection
+                ffmpeg = (
+                    FFmpeg()
+                    .option("y")
+                    .option("f", "lavfi")
+                    .option("i", "anullsrc")
+                    .option("f", "concat")
+                    .option("safe", 0)
+                    .option(
+                        "filter_complex",
+                        "colormatrix=bt470bg:bt709[0];"
+                        "[0]crop='iw-mod(iw,2)':'ih-mod(ih,2)'[main];"
+                        "[main]split[v1][v2];"
+                        "[v1]palettegen[pal];"
+                        "[v2][pal]paletteuse=dither=sierra2_4a",
+                    )
+                    .option("i", str(connect_config_file_path.resolve()))
+                    .option("pix_fmt", "yuv420p10le")
+                    .option("c:v", "libx265")
+                    .option("c:a", "aac")
+                    .option("crf", 0)
+                    .option("x265-params", "profile=main10")
+                    .option("shortest")
+                    .output(str(output_path))
+                )
+            else: # gif
+                output_path = directory / "out.gif"
+                ffmpeg = (
+                    FFmpeg()
+                    .option("y")
+                    .option("f", "concat")
+                    .option("safe", 0)
+                    .option("i", str(connect_config_file_path.resolve()))
+                    .option(
+                        "filter_complex",
+                        "colormatrix=bt470bg:bt709[main];"
+                        "[main]split[v1][v2];"
+                        "[v1]palettegen[pal];"
+                        "[v2][pal]paletteuse=dither=sierra2_4a",
+                    )
+                    .option("crf", 0)
+                    .output(str(output_path))
+                )
+            ffmpeg.on("completed", lambda: event.set())
+            await ffmpeg.execute()
+            await event.wait()
+
+            async with async_open(output_path, mode="rb") as file:
+                return await file.read()

src/async_pixiv/model/other/enums.py

@@ -1,12 +1,19 @@
 from async_pixiv.typedefs import Enum
+import enum
 
 __all__ = (
+    "Quality",
     "SearchTarget",
     "SearchShort",
     "SearchDuration",
     "SearchFilter",
 )
 
+class  Quality(Enum):
+    Square = enum.auto()
+    Medium = enum.auto()
+    Large = enum.auto()
+    Original = enum.auto()
 
 class SearchTarget(Enum):
     TAGS_PARTIAL = "partial_match_for_tags"  # 标签部分一致

src/async_pixiv/model/other/image.py

@@ -5,7 +5,7 @@
 from async_pixiv.model._base import PixivModel
 from async_pixiv.typedefs import URL
 
-__all__ = ("PixivImage", "ImageUrl", "AccountImage")
+__all__ = ("PixivImage", "QualityUrl", "ImageUrl", "AccountImage")
 
 
 class PixivImage(ABC, PixivModel):
@@ -17,8 +17,7 @@ def link(self) -> URL:
     async def download(self, *args, **kwargs):
         return await self.link.download(*args, **kwargs)
 
-
-class ImageUrl(PixivImage):
+class QualityUrl(PixivImage):
     square: URL | None = Field(None, alias="square_medium")
     medium: URL | None = None
     large: URL | None = None
@@ -28,6 +27,9 @@ class ImageUrl(PixivImage):
     def link(self) -> URL:
         return self.original or self.large or self.medium or self.square
 
+class ImageUrl(QualityUrl):
+    pass
+
 
 class AccountImage(PixivImage):
     small: URL | None = Field(None, alias="px_16x16")

src/async_pixiv/model/other/result.py

@@ -1,12 +1,13 @@
 from abc import ABC
 from typing import AsyncIterator, Iterator, Self
 
+from pydantic import Field
+
 from async_pixiv.model._base import PixivModel
+from async_pixiv.model.other.image import QualityUrl
 from async_pixiv.typedefs import URL
 from async_pixiv.utils.context import set_pixiv_client
 
-__all__ = ("PageResult",)
-
 
 class PageResult[T](ABC, PixivModel):
     previews: list[T] = []
@@ -30,3 +31,15 @@ async def iter_all_pages(self) -> AsyncIterator[T]:
             for result in next_results:
                 yield result
             next_results = await next_results.next()
+
+class UgoiraFrameMetadata(PixivModel):
+        file: str
+        delay: int
+
+
+class UgoiraMetadata(PixivModel):
+    zip_url: QualityUrl = Field(alias="zip_urls")
+    frames: list[UgoiraFrameMetadata]
+
+class UgoiraMetadataResult(PixivModel):
+    metadata: UgoiraMetadata = Field(alias="ugoira_metadata")

src/async_pixiv/model/other/tag.py

@@ -1,5 +1,7 @@
 from async_pixiv.model._base import PixivModel
 
+__all__ = ("Tag", "TagTranslation")
+
 
 class Tag(PixivModel):
     name: str

src/async_pixiv/utils/decorator.py

@@ -0,0 +1,32 @@
+import inspect
+from typing import Callable, ParamSpec
+
+from async_pixiv.utils.context import get_pixiv_client
+
+__all__ = ("auto_client",)
+
+
+P = ParamSpec("P")
+
+
+def auto_client[R](func: Callable[P, R]) -> Callable[P, R]:
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        signature = inspect.signature(func)
+        for name, param in signature.parameters.items():
+            if (
+                param.name == "client"
+                and any(
+                    [
+                        isinstance(param.annotation, str)
+                        and param.annotation == "PixivClient",
+                        isinstance(param.annotation, type)
+                        and param.annotation.__name__ == "PixivClient",
+                    ]
+                )
+                and kwargs.get("client") is None
+            ):
+                kwargs["client"] = get_pixiv_client()
+
+        return func(*args, **kwargs)
+
+    return wrapper

src/async_pixiv/utils/rate_limiter.py

@@ -0,0 +1,21 @@
+from typing import NoReturn
+
+from aiolimiter import AsyncLimiter
+
+__all__ = ("RateLimiter",)
+
+
+class RateLimiter(AsyncLimiter):
+    max_rate: float | None
+
+    def __init__(
+        self, max_rate: float | None = None, time_period: float = 60
+    ) -> NoReturn:
+        if max_rate is not None:
+            super().__init__(max_rate, time_period)
+        else:
+            self.max_rate = max_rate
+
+    async def acquire(self, amount: float = 1) -> None:
+        if self.max_rate is not None:
+            await super().acquire(amount)