TgStoGifBot/utils.py at master · SwaggyMacro/TgStoGifBot · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
"""
Utility functions for file uploads and ZIP management.
"""
import asyncio
import io
import os
import traceback

from loguru import logger
from telegram.error import RetryAfter, TimedOut


async def retry_upload_document(
    feedback_msg,
    file_obj,
    caption: str,
    parse_mode=None,
    max_retries: int = 5
) -> None:
    """
    Attempt to send a document (file path or BytesIO) with retry logic.

    :param feedback_msg: Telegram message object to reply to.
    :param file_obj: Either a file path (str) or a file-like object (e.g., BytesIO).
    :param caption: Caption text for the document.
    :param parse_mode: Parse mode for the caption (e.g., Markdown).
    :param max_retries: Maximum number of upload attempts.
    """
    for attempt in range(1, max_retries + 1):
        try:
            if isinstance(file_obj, str):
                # file_obj is a disk path
                with open(file_obj, "rb") as f:
                    await feedback_msg.reply_document(document=f, caption=caption, parse_mode=parse_mode)
            else:
                # file_obj is file-like (BytesIO)
                file_obj.seek(0)
                await feedback_msg.reply_document(document=file_obj, caption=caption, parse_mode=parse_mode)
            return  # successful upload
        except (RetryAfter, TimedOut) as e:
            logger.warning(f"Upload attempt {attempt} failed: {e}\n {traceback.format_exc()}")
            wait_secs = getattr(e, "retry_after", 5)
            await feedback_msg.reply_text(
                f"⚠️ Upload error: {e.__class__.__name__}. Retrying in {wait_secs}s (attempt {attempt}/{max_retries})…"
            )
            await asyncio.sleep(wait_secs)
        except Exception as e:
            await feedback_msg.reply_text(f"❌ Unexpected upload exception: `{e}`")
            raise

    await feedback_msg.reply_text(f"❌ Failed to upload after {max_retries} attempts.")


async def split_and_upload_document(
    feedback_msg,
    caption: str,
    zip_path: str,
    chunk_size: int = 50_000_000
) -> None:
    """
    Split a large ZIP into smaller parts and upload each.

    :param feedback_msg: Telegram message object to reply to.
    :param caption: Base caption for each part.
    :param zip_path: Path to the ZIP file.
    :param chunk_size: Maximum size (bytes) per part.
    """
    try:
        total_size = os.path.getsize(zip_path)
    except OSError as e:
        await feedback_msg.reply_text(f"❌ Could not stat '{zip_path}': {e}")
        return

    # If small enough, upload directly
    if total_size <= chunk_size:
        await retry_upload_document(feedback_msg, zip_path, caption, parse_mode="Markdown")
        return

    num_parts = (total_size + chunk_size - 1) // chunk_size
    await feedback_msg.reply_text(
        f"📦 Splitting large ZIP ({total_size // (1024*1024)} MB) into {num_parts} parts…"
    )

    base_name = os.path.basename(zip_path)
    windows_combine = ''
    linux_combine = ''
    macos_combine = ''

    with open(zip_path, "rb") as f:
        part_index = 1
        while True:
            chunk_data = f.read(chunk_size)
            if not chunk_data:
                break

            bio = io.BytesIO(chunk_data)
            bio.name = f"{base_name}.part{part_index:02d}"
            part_caption = f"{caption}\n(Part {part_index} of {num_parts})"
            if part_index == 1:
                windows_combine = f"```bash\ncopy /b {bio.name} + "
                linux_combine = f"```bash\ncat {bio.name} "
                macos_combine = f"```bash\ncat {bio.name} "
            else:
                windows_combine += f"{bio.name} + "
                linux_combine += f"{bio.name} "
                macos_combine += f"{bio.name} "
            await retry_upload_document(feedback_msg, bio, part_caption, parse_mode="Markdown")
            part_index += 1

    # Finalize combine commands
    windows_combine = windows_combine[:-3] + f" {base_name}.zip```"
    linux_combine += f"> {base_name}.zip```"
    macos_combine += f"> {base_name}.zip```"

    await feedback_msg.reply_text("✅ All parts uploaded successfully.\n"
                                  "🙋‍♂️ To combine them, use command below:\n"
                                  "🖥️ **Windows**\n"
                                  f"{windows_combine}\n"
                                  "🐧 **Linux**\n"
                                  f"{linux_combine}\n"
                                  "🍎 **macOS**\n"
                                  f"{macos_combine}\n\n"
                                  "**Make sure to run the command in the same directory where the parts are saved.**", parse_mode="Markdown")