moving to bytesio

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
                         )

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,26 +44,24 @@ 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
             width = int(video_stream['width']) or 0
             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'<code>{file.name}</code>'
+        )            
+        LOGGER.debug(f'Uploaded : {file.name}')
     if delete:
         os_remove(filepath)
 

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__

bot/plugins/split.py

@@ -9,53 +9,51 @@
 # 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)
     duration = int(float(video_stream["duration"]))
 
     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
+        LOGGER.debug(out_file)
+        yield out_file