From bc51ed109af2576b9063b9f4459a604281fac369 Mon Sep 17 00:00:00 2001 From: Davte Date: Mon, 22 Jun 2020 20:30:03 +0200 Subject: [PATCH] Split outgoing documents bigger than 50 MB and send chunks --- davtelepot/__init__.py | 2 +- davtelepot/bot.py | 92 +++++++++++++++++++++++++++++++++--------- davtelepot/messages.py | 6 ++- 3 files changed, 79 insertions(+), 21 deletions(-) diff --git a/davtelepot/__init__.py b/davtelepot/__init__.py index 0807d98..941f54f 100644 --- a/davtelepot/__init__.py +++ b/davtelepot/__init__.py @@ -11,7 +11,7 @@ __author__ = "Davide Testa" __email__ = "davide@davte.it" __credits__ = ["Marco Origlia", "Nick Lee @Nickoala"] __license__ = "GNU General Public License v3.0" -__version__ = "2.6.4" +__version__ = "2.6.5" __maintainer__ = "Davide Testa" __contact__ = "t.me/davte" diff --git a/davtelepot/bot.py b/davtelepot/bot.py index 8ec5f6f..31c57f5 100644 --- a/davtelepot/bot.py +++ b/davtelepot/bot.py @@ -97,6 +97,7 @@ class Bot(TelegramBot, ObjectWithDatabase, MultiLanguageObject): ] _log_file_name = None _errors_file_name = None + _documents_max_dimension = 50 * 1000 * 1000 # 50 MB def __init__( self, token, hostname='', certificate=None, max_connections=40, @@ -233,6 +234,7 @@ class Bot(TelegramBot, ObjectWithDatabase, MultiLanguageObject): self.shared_data = dict() self.Role = None self.packages = [sys.modules['davtelepot']] + self._documents_max_dimension = None # Add `users` table with its fields if missing if 'users' not in self.db.tables: table = self.db.create_table( @@ -594,6 +596,18 @@ class Bot(TelegramBot, ObjectWithDatabase, MultiLanguageObject): def administrators(self): return self._get_administrators(self) + @classmethod + def set_class_documents_max_dimension(cls, documents_max_dimension: int): + cls._documents_max_dimension = documents_max_dimension + + def set_documents_max_dimension(self, documents_max_dimension: int): + self._documents_max_dimension = documents_max_dimension + + @property + def documents_max_dimension(self) -> int: + return int(self._documents_max_dimension + or self.__class__._documents_max_dimension) + async def message_router(self, update, user_record, language): """Route Telegram `message` update to appropriate message handler.""" bot = self @@ -1796,25 +1810,26 @@ class Bot(TelegramBot, ObjectWithDatabase, MultiLanguageObject): # This buffered_file trick is necessary for two reasons # 1. File operations must be blocking, but sendDocument is a coroutine # 2. A `with` statement is not possible here - # `buffered_file` must be closed at all costs! - buffered_file = None if 'message' in update: update = update['message'] if chat_id is None and 'chat' in update: chat_id = self.get_chat_id(update) if reply_to_update and 'message_id' in update: reply_to_message_id = update['message_id'] - if ( - send_default_keyboard - and reply_markup is None - and type(chat_id) is int - and chat_id > 0 - and caption != self.authorization_denied_message - ): - reply_markup = self.get_keyboard( - update=update, - telegram_id=chat_id, - ) + if chat_id > 0: + user_record = self.db['users'].find_one(telegram_id=chat_id) + language = self.get_language(update=update, user_record=user_record) + if ( + send_default_keyboard + and reply_markup is None + and type(chat_id) is int + and caption != self.authorization_denied_message + ): + reply_markup = self.get_keyboard( + user_record=user_record + ) + else: + language = self.default_language if document_path is not None: with self.db as db: already_sent = db['sent_documents'].find_one( @@ -1839,16 +1854,58 @@ class Bot(TelegramBot, ObjectWithDatabase, MultiLanguageObject): ), 'rb' # Read bytes ) as file_: - buffered_file = io.BytesIO(file_.read()) - buffered_file.name = ( + file_size = os.fstat(file_.fileno()).st_size + document_chunks = ( + int( + file_size + / self.documents_max_dimension + ) + 1 + ) + original_document_name = ( document_name or file_.name or 'Document' ) - document = buffered_file + original_caption = caption + if '/' in original_document_name: + original_document_name = os.path.basename( + os.path.abspath(original_document_name) + ) + for i in range(document_chunks): + buffered_file = io.BytesIO( + file_.read(self.documents_max_dimension) + ) + if document_chunks > 1: + part = self.get_message( + 'davtelepot', 'part', + language=language + ) + caption = f"{original_caption} - {part} {i + 1}/{document_chunks}" + buffered_file.name = ( + f"{original_document_name} - " + f"{part} {i + 1}" + ) + else: + buffered_file.name = original_document_name + sent_document = await self.send_document( + chat_id=chat_id, + document=buffered_file, + thumb=thumb, + caption=caption, + parse_mode=parse_mode, + disable_notification=disable_notification, + reply_to_message_id=reply_to_message_id, + reply_markup=reply_markup, + update=update, + reply_to_update=reply_to_update, + send_default_keyboard=send_default_keyboard, + use_stored_file_id=use_stored_file_id + ) + return sent_document except FileNotFoundError as e: if buffered_file: buffered_file.close() + buffered_file = None return e else: use_stored_file_id = False @@ -1882,9 +1939,6 @@ class Bot(TelegramBot, ObjectWithDatabase, MultiLanguageObject): ), ['path'] ) - finally: - if buffered_file: - buffered_file.close() if ( type(sent_update) is dict and 'document' in sent_update diff --git a/davtelepot/messages.py b/davtelepot/messages.py index b38ecde..0df6c4d 100644 --- a/davtelepot/messages.py +++ b/davtelepot/messages.py @@ -14,7 +14,11 @@ davtelepot_messages = { "Questo messaggio รจ troppo lungo per essere inviato come " "messaggi separati.", } - } + }, + 'part': { + 'en': "part", + 'it': "parte", + }, } default_admin_messages = {