From 8aca932b682532f825614e33b8ab783c0b0c0b97 Mon Sep 17 00:00:00 2001 From: Aza Maulana Date: Mon, 24 Aug 2020 04:35:32 +0700 Subject: [PATCH] moving to bytesio --- bot/handlers/leech_handler.py | 2 + bot/handlers/upload_to_tg_handler.py | 79 ++++++++++++---------------- bot/plugins/IOHandler.py | 46 ++++++++++++++++ bot/plugins/split.py | 62 ++++++++++------------ 4 files changed, 111 insertions(+), 78 deletions(-) create mode 100644 bot/plugins/IOHandler.py diff --git a/bot/handlers/leech_handler.py b/bot/handlers/leech_handler.py index 3c8668e8..d54b10f4 100644 --- a/bot/handlers/leech_handler.py +++ b/bot/handlers/leech_handler.py @@ -64,6 +64,7 @@ async def func(client : Client, message: Message): for file in download.files: upload_status = await upload_to_tg_handler.func( os_path_join(dir, file.path), + client, reply, delete=True ) @@ -77,6 +78,7 @@ async def func(client : Client, message: Message): for file in download.files: upload_status = await upload_to_tg_handler.func( os_path_join(dir, file.path), + client, reply, delete=True ) diff --git a/bot/handlers/upload_to_tg_handler.py b/bot/handlers/upload_to_tg_handler.py index 7b5f687f..68513195 100644 --- a/bot/handlers/upload_to_tg_handler.py +++ b/bot/handlers/upload_to_tg_handler.py @@ -11,11 +11,11 @@ from os import path as os_path, listdir as os_lisdir, remove as os_remove, rmdir as os_rmdir from time import time from math import floor -from pyrogram import Message +from pyrogram import Client, Message from bot import LOCAL, CONFIG from bot.plugins import formater, split, thumbnail_video, ffprobe -async def func(filepath: str, message: Message, delete=False): +async def func(filepath: str, client: Client, message: Message, delete=False): if not os_path.exists(filepath): LOGGER.error(f'File not found : {filepath}') await message.edit_text( @@ -44,14 +44,11 @@ async def func(filepath: str, message: Message, delete=False): file_ext = os_path.splitext(filepath)[1].lower() LOGGER.debug(f'Uploading : {filepath}') - split_fn = None if file_ext in photo: - upload_fn = message.reply_photo - split_fn = split.func + upload_fn = client.send_photo elif file_ext in video: - split_fn = split.video - async def upload_fn(file, **kwargs): - probe = await ffprobe.func(file) + async def upload_fn(chat_id, file, **kwargs): + probe = await ffprobe.func(file.path) video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None) duration = int(float(video_stream["duration"])) or 0 @@ -59,11 +56,12 @@ async def upload_fn(file, **kwargs): height = int(video_stream['height']) or 0 await message.edit( LOCAL.GENERATE_THUMBNAIL.format( - name = os_path.basename(file) + name = file.name ) ) - thumbnail = await thumbnail_video.func(file) - await message.reply_video( + thumbnail = await thumbnail_video.func(file.path) + await client.send_video( + chat_id, file, supports_streaming=True, thumb=str(thumbnail) or None, @@ -75,8 +73,7 @@ async def upload_fn(file, **kwargs): if thumbnail: os_remove(str(thumbnail)) else: - upload_fn = message.reply_document - split_fn = split.func + upload_fn = client.send_document if os_path.getsize(filepath) > int(CONFIG.UPLOAD_MAX_SIZE): LOGGER.debug(f'File too large : {filepath}') @@ -85,40 +82,32 @@ async def upload_fn(file, **kwargs): name = os_path.basename(filepath) ) ) - splited = await split_fn(filepath, int(CONFIG.UPLOAD_MAX_SIZE)) - if not splited: - await message.edit( - LOCAL.SPLIT_FAILED.format( - name = os_path.basename(filepath) - ) - ) - return - for filepath in splited: - await message.edit( - LOCAL.UPLOADING_FILE.format( - name = os_path.basename(filepath) - ) + + async for file in split.func(filepath, int(CONFIG.UPLOAD_MAX_SIZE)): + await message.edit( + LOCAL.UPLOADING_FILE.format( + name = file.name ) - await func(filepath, message, delete=True) - return False + ) - info = { - "time" : time(), - "name" : os_path.basename(filepath), - "last_update" : 0, - "prev_text" : "" - } - await upload_fn( - filepath, - disable_notification=True, - progress=progress_upload_tg, - progress_args=( - message, - info - ), - caption=os_path.basename(filepath) - ) - LOGGER.debug(f'Uploaded : {filepath}') + info = { + "time" : time(), + "name" : file.name, + "last_update" : 0, + "prev_text" : "" + } + await upload_fn( + message.chat.id, + file, + disable_notification=True, + progress=progress_upload_tg, + progress_args=( + message, + info + ), + caption=f'{file.name}' + ) + LOGGER.debug(f'Uploaded : {file.name}') if delete: os_remove(filepath) diff --git a/bot/plugins/IOHandler.py b/bot/plugins/IOHandler.py new file mode 100644 index 00000000..40374e13 --- /dev/null +++ b/bot/plugins/IOHandler.py @@ -0,0 +1,46 @@ +# GOAL: +# getting track for logging + +import logging + +LOGGER = logging.getLogger(__name__) + +# GOAL: +# handling file upload using bytesio + +from io import FileIO +from os import path as os_path, remove as os_remove + +class ChunkIO(FileIO): + + def __init__(self, filepath, pos=0, size=-1): + FileIO.__init__(self, filepath, 'rb') + self.name = os_path.basename(filepath) + self.__filepath__ = filepath + self.__startpos__ = pos + self.__size__ = size + self.__currentpos__ = 0 + + total_size = os_path.getsize(filepath) - pos + if self.__size__ > total_size: + self.__size__ = total_size + + + def read(self, size = -1): + if size == -1 or size + self.__currentpos__ > self.__size__: + size = self.__size__ - self.__currentpos__ + self.__currentpos__ += size + return FileIO.read(self, size) + + def seek(self, position, whence=0): + if whence == 1: + position += self.__currentpos__ + elif whence == 2: + position += self.__size__ + + self.__currentpos__ = position + FileIO.seek(self, self.__startpos__ + position) + return position + + def tell(self): + return self.__currentpos__ diff --git a/bot/plugins/split.py b/bot/plugins/split.py index d097bf7a..c0065f28 100644 --- a/bot/plugins/split.py +++ b/bot/plugins/split.py @@ -9,45 +9,44 @@ # create split handler class from os import path as os_path, remove as os_remove +from io import FileIO import asyncio -from glob import glob -from bot.plugins import ffprobe +from bot.plugins import ffprobe, IOHandler async def func(filepath, size): - if not os_path.isfile(filepath): - LOGGER.error('File not found : ' + filepath) - return False - - cmd = [ - "split", - "--numeric-suffixes=1", - "--suffix-length=3", - f"--bytes={size}", - filepath, - filepath + "." - ] - LOGGER.debug(cmd) - - process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - await process.communicate() - - list = glob(filepath + ".*").sort() - LOGGER.debug(list) - return list + file_ext = os_path.splitext(filepath)[1] + video_units = ['.mp4','.mkv','.avi','.webm','.wmv','.mov'] + if file_ext in video_units: + async for splitted_video in video(filepath, size): + with FileIO(splitted_video, 'rb') as f: + f.name = os_path.basename(splitted_video) + f.path = splitted_video + yield f + os_remove(splitted_video) + else: + total_size = os_path.getsize(filepath) + pos = 0 + index = 0 + while pos < total_size: + index += 1 + with IOHandler.ChunkIO(filepath, pos, size) as f: + prefix = '' + if size < total_size: + prefix = f'.{index:03d}' + pos += size + f.name = os_path.basename(filepath) + prefix + yield f async def video(filepath, size): supported = ['.mp4','.mkv','.avi','.webm','.wmv','.mov'] if not os_path.isfile(filepath): LOGGER.error('File not found : ' + filepath) - return False + raise Exception('File not found') file_path_name, file_ext = os_path.splitext(filepath) if not file_ext in supported: - return False + LOGGER.error('File not supported : ' + filepath) + raise Exception('File not supported') probe = await ffprobe.func(filepath) video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None) @@ -55,7 +54,6 @@ async def video(filepath, size): splited_duration = 0 i = 0 - list = [] while splited_duration < duration: i+=1 out_file = file_path_name + ".{:03d}".format(i) + file_ext @@ -89,7 +87,5 @@ async def video(filepath, size): splited_duration += int(float(video_stream["duration"])) - list.append(out_file) - - LOGGER.debug(list) - return list \ No newline at end of file + LOGGER.debug(out_file) + yield out_file \ No newline at end of file