From 992493aa15c0d0b912af9d32a95a31c57c0458ac Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Fri, 3 Nov 2023 15:57:29 +0300 Subject: [PATCH] recover file url if not present --- e621.py | 3 +- main.py | 125 +++++++++++++++++++++++++++++--------------------------- 2 files changed, 65 insertions(+), 63 deletions(-) diff --git a/e621.py b/e621.py index 10a32ae..8c6f671 100644 --- a/e621.py +++ b/e621.py @@ -136,9 +136,8 @@ 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=100) -> List[E621Post]: + 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: diff --git a/main.py b/main.py index 2d2f596..8ae9849 100644 --- a/main.py +++ b/main.py @@ -25,7 +25,7 @@ from aiogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton, B CallbackQuery import dotenv -from e621 import E621, E621Post +from e621 import E621, E621Post, E621PostFile dotenv.load_dotenv('.env') @@ -43,6 +43,10 @@ def format_tags(tags: Iterable[str]): return ' '.join(f'#{tag}' for tag in tags) or 'None' +def recover_url(file: E621PostFile): + return f'https://static1.e621.net/data/{file.md5[:2]}/{file.md5[:2:4]}/{file.md5}.{file.ext}' + + async def send_post(post: E621Post, tag_list: List[str]): try: logging.warning(f'Sending post #{post.id}') @@ -58,68 +62,67 @@ async def send_post(post: E621Post, tag_list: List[str]): copyright_tags and f'Copyright: {format_tags(copyright_tags)}', f'\nhttps://e621.net/posts/{post.id}' ] if l) - if post.file.url: - try: - logging.warning(post.file.url) - async with httpx.AsyncClient() as client: + if not post.file.url: + post.file.url = recover_url(post.file) + try: + logging.warning(post.file.url) + async with httpx.AsyncClient() as client: + file = BytesIO() + file.write((await client.get(post.file.url)).content) + file.name = f'file.{post.file.ext}' + file.seek(0) + if post.file.ext in ('webm', 'gif'): + with TemporaryDirectory() as td: + src_path = Path(td) / f'video.{post.file.ext}' + mp4_path = Path(td) / 'video.mp4' + with open(src_path, 'wb') as webm: + webm.write(file.read()) + video_input = ffmpeg\ + .input(str(src_path)) + cmd = video_input \ + .output(str(mp4_path), + vf='pad=width=ceil(iw/2)*2:height=ceil(ih/2)*2:x=0:y=0:color=Black', + vcodec='libx264', + crf='26') + logging.info('ffmpeg ' + ' '.join(cmd.get_args())) + cmd.run() + s3 = boto3.client('s3', aws_access_key_id=os.environ['AWS_ACCESS_KEY'], aws_secret_access_key=os.environ['AWS_SECRET_KEY']) + bucket = os.environ['AWS_S3_BUCKET'] + upload_filename = f'e621-{post.id}-{int(time())}.mp4' + s3.upload_file(mp4_path, bucket, upload_filename, ExtraArgs={'ACL': 'public-read', 'ContentType': 'video/mp4'}) + await bot.send_message(int(os.environ['SEND_CHAT']), + f'https://{bucket}.s3.amazonaws.com/{upload_filename}\n\n' + caption, + parse_mode=ParseMode.HTML) + src_path.unlink() + mp4_path.unlink() + elif post.file.ext in ('png', 'jpg'): + markup = InlineKeyboardMarkup(inline_keyboard=[[ + InlineKeyboardButton(text='NSFW', callback_data='send nsfw'), + InlineKeyboardButton(text='Safe', callback_data='send pics'), + ]]) + # if post.file.size > 10000000: + logging.warning('compressing') + dl_im = Image.open(file).convert('RGBA') + size = dl_im.size + if size[0] > 2000 or size[1] > 2000: + larger_dimension = max(size) + ratio = 2000 / larger_dimension + dl_im = dl_im.resize((int(size[0] * ratio), int(size[1] * ratio)), + Image.LANCZOS) + print(f'Resizing from {size[0]}x{size[1]} to {dl_im.size[0]}x{dl_im.size[1]}') + im = Image.new('RGBA', dl_im.size, (255, 255, 255)) + composite = Image.alpha_composite(im, dl_im).convert('RGB') file = BytesIO() - file.write((await client.get(post.file.url)).content) - file.name = f'file.{post.file.ext}' + composite.save(file, format='JPEG') file.seek(0) - if post.file.ext in ('webm', 'gif'): - with TemporaryDirectory() as td: - src_path = Path(td) / f'video.{post.file.ext}' - mp4_path = Path(td) / 'video.mp4' - with open(src_path, 'wb') as webm: - webm.write(file.read()) - video_input = ffmpeg\ - .input(str(src_path)) - cmd = video_input \ - .output(str(mp4_path), - vf='pad=width=ceil(iw/2)*2:height=ceil(ih/2)*2:x=0:y=0:color=Black', - vcodec='libx264', - crf='26') - logging.info('ffmpeg ' + ' '.join(cmd.get_args())) - cmd.run() - s3 = boto3.client('s3', aws_access_key_id=os.environ['AWS_ACCESS_KEY'], aws_secret_access_key=os.environ['AWS_SECRET_KEY']) - bucket = os.environ['AWS_S3_BUCKET'] - upload_filename = f'e621-{post.id}-{int(time())}.mp4' - s3.upload_file(mp4_path, bucket, upload_filename, ExtraArgs={'ACL': 'public-read', 'ContentType': 'video/mp4'}) - await bot.send_message(int(os.environ['SEND_CHAT']), - f'https://{bucket}.s3.amazonaws.com/{upload_filename}\n\n' + caption, - parse_mode=ParseMode.HTML) - src_path.unlink() - mp4_path.unlink() - elif post.file.ext in ('png', 'jpg'): - markup = InlineKeyboardMarkup(inline_keyboard=[[ - InlineKeyboardButton(text='NSFW', callback_data='send nsfw'), - InlineKeyboardButton(text='Safe', callback_data='send pics'), - ]]) - # if post.file.size > 10000000: - logging.warning('compressing') - dl_im = Image.open(file).convert('RGBA') - size = dl_im.size - if size[0] > 2000 or size[1] > 2000: - larger_dimension = max(size) - ratio = 2000 / larger_dimension - dl_im = dl_im.resize((int(size[0] * ratio), int(size[1] * ratio)), - Image.LANCZOS) - print(f'Resizing from {size[0]}x{size[1]} to {dl_im.size[0]}x{dl_im.size[1]}') - im = Image.new('RGBA', dl_im.size, (255, 255, 255)) - composite = Image.alpha_composite(im, dl_im).convert('RGB') - file = BytesIO() - composite.save(file, format='JPEG') - file.seek(0) - await bot.send_photo(int(os.environ['SEND_CHAT']), - BufferedInputFile(file.read(), 'file.jpg'), - caption=caption, - parse_mode=ParseMode.HTML, - reply_markup=markup) - await redis.sadd('e621:sent', post.id) - except Exception as e: - logging.exception(e) - else: - logging.warning('Not sending because url is None', post.file) + await bot.send_photo(int(os.environ['SEND_CHAT']), + BufferedInputFile(file.read(), 'file.jpg'), + caption=caption, + parse_mode=ParseMode.HTML, + reply_markup=markup) + await redis.sadd('e621:sent', post.id) + except Exception as e: + logging.exception(e) except Exception as e: logging.exception(e)