diff --git a/classes/pyroclient.py b/classes/pyroclient.py index 0d03584..38b900c 100644 --- a/classes/pyroclient.py +++ b/classes/pyroclient.py @@ -8,11 +8,13 @@ from shutil import rmtree from time import time from traceback import format_exc from typing import List, Tuple, Union +import aiofiles import pyrogram from aiohttp import ClientSession from bson import ObjectId from dateutil.relativedelta import relativedelta +from classes.enums.submission_types import SubmissionType from libbot import json_read, json_write from libbot.i18n import BotLocale from libbot.i18n.sync import _ @@ -43,10 +45,14 @@ from classes.exceptions import ( ) from classes.pyrocommand import PyroCommand from modules.api_client import ( - BodyPhotoUploadAlbumsAlbumPhotosPost, + BodyPhotoUpload, + BodyVideoUpload, File, + Photo, + Video, client, photo_upload, + video_upload, ) from modules.database import col_submitted from modules.http_client import http_session @@ -398,7 +404,7 @@ class PyroClient(Client): language_code=command_set.language_code, ) - async def submit_photo( + async def submit_media( self, id: str ) -> Tuple[Union[Message, None], Union[str, None]]: db_entry = col_submitted.find_one({"_id": ObjectId(id)}) @@ -431,24 +437,47 @@ class PyroClient(Client): db_entry["user"], db_entry["telegram"]["msg_id"] ) - with open(str(filepath), "rb") as fh: - photo_bytes = BytesIO(fh.read()) + async with aiofiles.open(str(filepath), "rb") as fh: + media_bytes = BytesIO(await fh.read()) try: - response = await photo_upload( - self.config["posting"]["api"]["album"], - client=client, - multipart_data=BodyPhotoUploadAlbumsAlbumPhotosPost( - File(photo_bytes, filepath.name, "image/jpeg") - ), - ignore_duplicates=self.config["submission"]["allow_duplicates"], - compress=False, - caption="queue", - ) + if db_entry["type"] == SubmissionType.PHOTO.value: + response = await photo_upload( + self.config["posting"]["api"]["album"], + client=client, + multipart_data=BodyPhotoUpload( + File(media_bytes, filepath.name, "image/jpeg") + ), + ignore_duplicates=self.config["submission"]["allow_duplicates"], + compress=False, + caption="queue", + ) + elif db_entry["type"] == SubmissionType.VIDEO.value: + response = await video_upload( + self.config["posting"]["api"]["album"], + client=client, + multipart_data=BodyVideoUpload( + File(media_bytes, filepath.name, "video/*") + ), + caption="queue", + ) + elif db_entry["type"] == SubmissionType.ANIMATION.value: + response = await video_upload( + self.config["posting"]["api"]["album"], + client=client, + multipart_data=BodyVideoUpload( + File(media_bytes, filepath.name, "video/*") + ), + caption="queue", + ) except UnexpectedStatus: raise SubmissionUnsupportedError(str(filepath)) - response_dict = loads(response.content.decode("utf-8")) + response_dict = ( + {} + if not hasattr(response, "content") + else loads(response.content.decode("utf-8")) + ) if "duplicates" in response_dict and len(response_dict["duplicates"]) > 0: duplicates = [] @@ -480,7 +509,10 @@ class PyroClient(Client): except (FileNotFoundError, NotADirectoryError): logger.error("Could not delete '%s' on submission accepted", filepath) - return submission, response.parsed.id + return ( + submission, + response.id if not hasattr(response, "parsed") else response.parsed.id, + ) async def ban_user(self, id: int) -> None: pass diff --git a/config_example.json b/config_example.json index 8c9fded..0d0f248 100644 --- a/config_example.json +++ b/config_example.json @@ -56,6 +56,11 @@ "ignore_admins": true, "text": "#submitted" }, + "types": { + "photo": true, + "video": false, + "animation": false + }, "extensions": { "photo": [ "jpg", diff --git a/modules/api_client.py b/modules/api_client.py index 515f28d..8192499 100644 --- a/modules/api_client.py +++ b/modules/api_client.py @@ -35,14 +35,26 @@ from photosapi_client.api.default.user_me_users_me_get import sync as user_me from photosapi_client.api.default.video_find_albums_album_videos_get import ( asyncio as video_find, ) +from photosapi_client.api.default.video_get_videos_id_get import asyncio as video_get +from photosapi_client.api.default.video_patch_videos_id_patch import ( + asyncio as video_patch, +) +from photosapi_client.api.default.video_upload_albums_album_videos_post import ( + asyncio as video_upload, +) from photosapi_client.models.body_login_for_access_token_token_post import ( BodyLoginForAccessTokenTokenPost, ) from photosapi_client.models.body_photo_upload_albums_album_photos_post import ( - BodyPhotoUploadAlbumsAlbumPhotosPost, + BodyPhotoUploadAlbumsAlbumPhotosPost as BodyPhotoUpload, +) +from photosapi_client.models.body_video_upload_albums_album_videos_post import ( + BodyVideoUploadAlbumsAlbumVideosPost as BodyVideoUpload, ) from photosapi_client.models.http_validation_error import HTTPValidationError +from photosapi_client.models.photo import Photo from photosapi_client.models.token import Token +from photosapi_client.models.video import Video from photosapi_client.types import File from modules.http_client import http_session diff --git a/modules/sender.py b/modules/sender.py index 3456446..ca1db19 100644 --- a/modules/sender.py +++ b/modules/sender.py @@ -3,7 +3,7 @@ from datetime import datetime from os import makedirs, path from random import choice from shutil import rmtree -from traceback import format_exc +from traceback import format_exc, print_exc from uuid import uuid4 import aiofiles @@ -12,7 +12,16 @@ from photosapi_client.errors import UnexpectedStatus from PIL import Image from pyrogram.client import Client -from modules.api_client import authorize, client, photo_find, photo_patch +from modules.api_client import ( + authorize, + client, + photo_find, + photo_get, + photo_patch, + video_find, + video_get, + video_patch, +) from modules.database import col_sent, col_submitted logger = logging.getLogger(__name__) @@ -30,9 +39,19 @@ async def send_content(app: Client, http_session: ClientSession) -> None: return try: - pic = choice( + funcs = [] + if app.config["posting"]["types"]["photo"]: + funcs.append(photo_find) + if ( + app.config["posting"]["types"]["video"] + or app.config["posting"]["types"]["animation"] + ): + funcs.append(video_find) + func = choice(funcs) + + media = choice( ( - await photo_find( + await func( album=app.config["posting"]["api"]["album"], caption="queue", page_size=app.config["posting"]["page_size"], @@ -56,41 +75,48 @@ async def send_content(app: Client, http_session: ClientSession) -> None: ) return - response = await http_session.get( - f"{app.config['posting']['api']['address']}/photos/{pic.id}", - headers={"Authorization": f"Bearer {token}"}, - ) - - if response.status != 200: + try: + if func is photo_find: + response = await photo_get(id=media.id, client=client) + else: + response = await video_get(id=media.id, client=client) + except Exception as exp: + print_exc() logger.warning( - app._("post_invalid_pic", "console").format( - response.status, str(await response.json()) - ) + "Media is invalid: %s", + exp + # app._("post_invalid_pic", "console").format( + # response.status, str(await response.json()) + # ) ) if app.config["reports"]["error"]: - await app.send_message( - app.owner, - app._("post_invalid_pic", "message").format( - response.status, await response.json() - ), - ) + await app.send_message(app.owner, f"Media is invalid: {exp}") + return + # await app.send_message( + # app.owner, + # app._("post_invalid_pic", "message").format( + # response.status, await response.json() + # ), + # ) tmp_dir = str(uuid4()) makedirs(path.join(app.config["locations"]["tmp"], tmp_dir), exist_ok=True) - tmp_path = path.join(tmp_dir, pic.filename) + tmp_path = path.join(tmp_dir, media.filename) async with aiofiles.open( path.join(app.config["locations"]["tmp"], tmp_path), "wb" ) as out_file: - await out_file.write(await response.read()) + await out_file.write(response.payload.read()) logger.info( - f"Candidate {pic.filename} ({pic.id}) is {path.getsize(path.join(app.config['locations']['tmp'], tmp_path))} bytes big", + f"Candidate {media.filename} ({media.id}) is {path.getsize(path.join(app.config['locations']['tmp'], tmp_path))} bytes big", ) - if path.getsize(path.join(app.config["locations"]["tmp"], tmp_path)) > 5242880: + if ( + path.getsize(path.join(app.config["locations"]["tmp"], tmp_path)) > 5242880 + ) and func is photo_find: image = Image.open(path.join(app.config["locations"]["tmp"], tmp_path)) width, height = image.size image = image.resize((int(width / 2), int(height / 2)), Image.ANTIALIAS) @@ -110,7 +136,9 @@ async def send_content(app: Client, http_session: ClientSession) -> None: ) image.close() - if path.getsize(path.join(app.config["locations"]["tmp"], tmp_path)) > 5242880: + if ( + path.getsize(path.join(app.config["locations"]["tmp"], tmp_path)) > 5242880 + ) and func is photo_find: rmtree( path.join(app.config["locations"]["tmp"], tmp_dir), ignore_errors=True ) @@ -118,7 +146,7 @@ async def send_content(app: Client, http_session: ClientSession) -> None: del response - submitted = col_submitted.find_one({"temp.file": pic.filename}) + submitted = col_submitted.find_one({"temp.file": media.filename}) if submitted is not None and submitted["caption"] is not None: caption = submitted["caption"].strip() @@ -150,14 +178,24 @@ async def send_content(app: Client, http_session: ClientSession) -> None: caption = caption try: - sent = await app.send_photo( - app.config["posting"]["channel"], - path.join(app.config["locations"]["tmp"], tmp_path), - caption=caption, - disable_notification=app.config["posting"]["silent"], - ) + if func is photo_find: + sent = await app.send_photo( + app.config["posting"]["channel"], + path.join(app.config["locations"]["tmp"], tmp_path), + caption=caption, + disable_notification=app.config["posting"]["silent"], + ) + else: + sent = await app.send_video( + app.config["posting"]["channel"], + path.join(app.config["locations"]["tmp"], tmp_path), + caption=caption, + disable_notification=app.config["posting"]["silent"], + ) except Exception as exp: - logger.error(f"Could not send image {pic.filename} ({pic.id}) due to {exp}") + logger.error( + f"Could not send media {media.filename} ({media.id}) due to {exp}" + ) if app.config["reports"]["error"]: await app.send_message( app.owner, @@ -169,8 +207,8 @@ async def send_content(app: Client, http_session: ClientSession) -> None: col_sent.insert_one( { "date": datetime.now(), - "image": pic.id, - "filename": pic.filename, + "image": media.id, + "filename": media.filename, "channel": app.config["posting"]["channel"], "caption": None if (submitted is None or submitted["caption"] is None) @@ -178,13 +216,14 @@ async def send_content(app: Client, http_session: ClientSession) -> None: } ) - await photo_patch(id=pic.id, client=client, caption="sent") + func_patch = photo_patch if func is photo_find else video_patch + await func_patch(id=media.id, client=client, caption="sent") rmtree(path.join(app.config["locations"]["tmp"], tmp_dir), ignore_errors=True) logger.info( app._("post_sent", "console").format( - pic.id, + media.id, str(app.config["posting"]["channel"]), caption.replace("\n", "%n"), str(app.config["posting"]["silent"]), diff --git a/plugins/commands/photos.py b/plugins/commands/photos.py index 8e9123b..9a17fad 100644 --- a/plugins/commands/photos.py +++ b/plugins/commands/photos.py @@ -18,7 +18,7 @@ from ujson import loads from classes.pyroclient import PyroClient from modules.api_client import ( - BodyPhotoUploadAlbumsAlbumPhotosPost, + BodyPhotoUpload, File, client, photo_delete, @@ -158,7 +158,7 @@ async def cmd_import(app: PyroClient, msg: Message): uploaded = await photo_upload( app.config["posting"]["api"]["album"], client=client, - multipart_data=BodyPhotoUploadAlbumsAlbumPhotosPost( + multipart_data=BodyPhotoUpload( File(photo_bytes, Path(filename).name, "image/jpeg") ), ignore_duplicates=app.config["submission"]["allow_duplicates"], diff --git a/plugins/handlers/submission.py b/plugins/handlers/submission.py index b184657..a7b1faf 100644 --- a/plugins/handlers/submission.py +++ b/plugins/handlers/submission.py @@ -215,7 +215,7 @@ async def get_submission(app: PyroClient, msg: Message): and app.config["submission"]["require_confirmation"]["admins"] is False ): try: - submitted = await app.submit_photo(str(inserted.inserted_id)) + submitted = await app.submit_media(str(inserted.inserted_id)) await msg.reply_text( app._("sub_yes_auto", "message", locale=user_locale), disable_notification=True,