Feat: possible to download gif.

awesome/plugins/setu/__init__.py

@@ -1,5 +1,4 @@
-from os import getcwd, path
-from os.path import join
+from os import path
 from random import choice, randint
 from re import split, findall
 from time import time
@@ -23,9 +22,10 @@
 from Services.util.sauce_nao_helper import sauce_helper
 from awesome.Constants import user_permission as perm, group_permission
 from awesome.Constants.function_key import SETU, TRIGGER_BLACKLIST_WORD, HIT_XP
-from awesome.Constants.path_constants import DL_PATH
+from awesome.Constants.path_constants import DL_PATH, PIXIV_PIC_PATH
 from awesome.Constants.plugins_command_constants import PROMPT_FOR_KEYWORD
 from awesome.adminControl import setu_function_control, get_privilege, group_control, user_control
+from awesome.plugins.setu.setu_utilties import download_gif
 from awesome.plugins.setu.setuconfig import SetuConfig
 from config import SUPER_USER, PIXIV_REFRESH_TOKEN
 from util.helper_util import anime_reverse_search_response, set_group_permission, construct_message_chain
@@ -554,6 +554,19 @@ def _get_image_data_from_username(key_word: str) -> (str, str):
 
 
 async def _download_pixiv_image_helper(illust) -> str:
+    if illust.type != 'ugoira':
+        return await _handle_normal_illust_download(illust)
+
+    ugoira_data = pixiv_api.ugoira_metadata(illust.id)
+    url_list = ugoira_data.ugoira_metadata.zip_urls.medium
+    duration = ugoira_data.ugoira_metadata.frames[0].delay
+    gif_path = await download_gif(url_list, illust.user.name + '_' + illust.title, duration)
+
+    logger.success(f'Gif path done: {gif_path}')
+    return gif_path if gif_path else ''
+
+
+async def _handle_normal_illust_download(illust):
     if illust['meta_single_page']:
         if 'original_image_url' in illust['meta_single_page']:
             image_url = illust.meta_single_page['original_image_url']
@@ -569,19 +582,15 @@ async def _download_pixiv_image_helper(illust) -> str:
             image_urls['large'] if 'large' in image_urls else \
                 image_urls['medium'] if 'medium' in image_urls else \
                     image_urls['square_medium']
-
     logger.info(f"{illust.title}: {image_url}, {illust.id}")
-    setu_file_path = join(getcwd(), 'data', 'pixivPic')
-
+    setu_file_path = PIXIV_PIC_PATH
     try:
         setu_file_path = await download_image(
             image_url, setu_file_path, headers={'Referer': 'https://app-api.pixiv.net/'})
     except Exception as err:
         logger.info(f'Download image error: {err}')
         return ''
-
     edited_path = await slight_adjust_pic_and_get_path(setu_file_path)
-
     logger.info("PATH = " + edited_path)
     return edited_path
 

awesome/plugins/setu/setu_utilties.py

@@ -0,0 +1,65 @@
+from os import path, mkdir, remove, rmdir, walk
+from os.path import exists, abspath, join
+from typing import Optional, List
+from zipfile import ZipFile
+
+from PIL import Image
+from nonebot import logger
+from youtube_dl.utils import sanitize_filename
+
+from Services.util import global_httpx_client
+from awesome.Constants.path_constants import PIXIV_PIC_PATH
+
+
+async def download_gif(ugoira_url: str, title: str, duration: float) -> Optional[str]:
+    if ugoira_url is None:
+        logger.warning('Zip ugoira is None!!!')
+        return None
+
+    title = sanitize_filename(title)
+    gif_zip_path = path.join(PIXIV_PIC_PATH, title).__str__()
+    zip_file_name = ugoira_url.split('/')[-1]
+
+    return await _download_gif_processor(zip_file_name, ugoira_url, gif_zip_path, duration)
+
+
+async def _download_gif_processor(
+        zip_file_name: str, ugoira_url: str, gif_zip_path: str, duration: float) -> Optional[str]:
+    gif_zip_path_path = None
+    zip_file_path = path.join(PIXIV_PIC_PATH, zip_file_name)
+    if not exists(gif_zip_path):
+        zip_file_path = \
+            await global_httpx_client.download(
+                ugoira_url, zip_file_path, headers={'Referer': 'https://app-api.pixiv.net/'})
+
+        gif_zip_path_path = gif_zip_path.split('.gif')[0] + '.gif'
+        gif_zip_path = gif_zip_path.split('.')[0]
+        with ZipFile(zip_file_path, 'r') as zipref:
+            if not exists(gif_zip_path):
+                mkdir(gif_zip_path)
+            zipref.extractall(gif_zip_path)
+            im = Image.open(_get_absolute_file_paths(gif_zip_path)[0])
+            logger.info('Making the gif...')
+            im.save(f'{gif_zip_path_path}', save_all=True,
+                    append_images=[Image.open(file) for file in _get_absolute_file_paths(gif_zip_path)],
+                    duration=duration,
+                    loop=0)
+
+        logger.info('Removing gif making single img cache...')
+        remove(zip_file_path)
+        for file in _get_absolute_file_paths(gif_zip_path):
+            remove(file)
+
+        logger.info('Removing zip cache.')
+        rmdir(gif_zip_path)
+
+    return gif_zip_path_path
+
+
+def _get_absolute_file_paths(directory) -> List[str]:
+    file_paths = []
+    for folder, subs, files in walk(directory):
+        for filename in files:
+            file_paths.append(abspath(join(folder, filename)))
+
+    return file_paths