Commit b3eba16

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-01-31 02:04:20
refactor(sender): refactor `send2tg` function
1 parent 8845ccc
Changed files (1)
src
messages
src/messages/sender.py
@@ -124,6 +124,64 @@ def preprocess_media(media: list[dict]) -> list[dict]:
     return results
 
 
+async def send_texts(
+    client: Client,
+    target_chat: int | str,
+    reply_parameters: ReplyParameters,
+    *,
+    texts: str = "",
+    cooldown: float = 0,
+) -> list[Message | None]:
+    sent_messages: list[Message | None] = []
+    logger.trace(f"Sending {len(texts)} texts only")
+    for idx, msg in enumerate(smart_split(texts.strip())):
+        if not msg:
+            continue
+        if idx == 0:
+            sent_messages.append(await client.send_message(target_chat, msg, reply_parameters=reply_parameters))
+        else:  # disbale reply
+            sent_messages.append(await client.send_message(target_chat, msg, reply_parameters=ReplyParameters()))
+            await asyncio.sleep(cooldown)
+    return sent_messages
+
+
+async def send_single_media(
+    client: Client,
+    target_chat: int | str,
+    reply_parameters: ReplyParameters,
+    *,
+    media: dict,
+    texts: str = "",
+    cooldown: float = 0,
+    **kwargs,
+) -> list[Message | None]:
+    sent_messages: list[Message | None] = []
+    logger.trace(f"Sending single media with {len(texts)} texts")
+    caption = smart_split(texts, CAPTION_LENGTH)[0]
+    remaining_texts = texts.removeprefix(caption)
+    if photo := media.get("photo"):
+        sent_messages.append(await client.send_photo(chat_id=target_chat, photo=photo, caption=caption, reply_parameters=reply_parameters))
+    elif video := media.get("video"):
+        sent_messages.append(
+            await client.send_video(
+                chat_id=target_chat,
+                reply_parameters=reply_parameters,
+                caption=caption,
+                progress=telegram_uploading,
+                progress_args=(kwargs.get("progress", False), video, kwargs.get("detail_progress", True)),
+                **media,
+            )
+        )
+    if remaining_texts:
+        sent_messages.extend(await send_texts(client, target_chat, ReplyParameters(), texts=remaining_texts, cooldown=cooldown))
+
+    for key in ["path", "thumb", "audio", "photo", "video"]:
+        if media.get(key) and Path(media[key]).is_file():
+            logger.trace(f"Deleting: {media[key]}")
+            Path(media[key]).unlink(missing_ok=True)
+    return sent_messages
+
+
 async def send2tg(
     client: Client,
     message: Message,
@@ -169,19 +227,14 @@ async def send2tg(
             "video": "path/to/video.mp4",
         }
     ]
-    TODO: Support to send audio and document
     """
     if not target_chat:
         target_chat = kwargs["target_chat"] if kwargs.get("target_chat") else message.chat.id
     target_chat = to_int(target_chat)
     reply_parameters = get_reply_to(message.id, reply_msg_id)
-    sent_messages: list[Message | None] = []  # save sent messages results
-
-    if media is None:
-        media = []
+    media = media or []
     media = preprocess_media(media)
-    if comments is None:
-        comments = []
+    comments = comments or []
 
     # no text, but has comments. treat comments as texts
     texts = texts if texts else "".join(comments).strip()
@@ -192,34 +245,6 @@ async def send2tg(
     if kwargs.get("progress") and len(media) > 0:
         await modify_progress(text=f"⏫正在上传:\n{summay_media(media)}", force_update=True, **kwargs)
 
-    # only media
-    if media and not texts:
-        logger.trace(f"Sending {len(media)} media without any texts")
-        if len(media) == 1:
-            if media[0].get("photo"):
-                sent_messages.append(await client.send_photo(chat_id=target_chat, photo=media[0]["photo"], reply_parameters=reply_parameters))
-            elif media[0].get("video"):
-                sent_messages.append(
-                    await client.send_video(
-                        chat_id=target_chat,
-                        reply_parameters=reply_parameters,
-                        progress=telegram_uploading,
-                        progress_args=(kwargs.get("progress", False), media[0]["video"], kwargs.get("detail_progress", True)),
-                        **media[0],
-                    )
-                )
-        elif 1 < len(media) <= 10:
-            group = warp_media_group(media)
-            sent_messages.extend(await client.send_media_group(target_chat, media=group, reply_parameters=reply_parameters))
-        else:
-            media_chunks = [media[i : i + 10] for i in range(0, len(media), 10)]
-            for idx, chunk in enumerate(media_chunks):
-                if idx == 0:
-                    sent_messages.extend(await send2tg(client, message, target_chat, reply_msg_id, media=chunk, **kwargs))
-                else:
-                    sent_messages.extend(await send2tg(client, message, target_chat, reply_msg_id=-1, media=chunk, **kwargs))  # disbale reply
-                await asyncio.sleep(cooldown)  # cool down
-
     # append comments to texts
     # For len(texts) < 1024 , ensure the combined texts and comments remains below 1024 characters to avoid sending a subsequent message containing only the comments.
     # For long texts, keep all comments
@@ -230,103 +255,35 @@ async def send2tg(
     else:
         texts = texts + "".join(comments)
 
-    videos = [x for x in media if x.get("video")]
-    photos = [x for x in media if x.get("photo")]
-    logger.trace(f"{len(texts)} texts, {len(comments)} comments, {len(videos)} videos, {len(photos)} photos: {texts!r}")
-
-    # only texts
-    if texts and not media:
-        logger.trace(f"Sending {len(texts)} texts without any media")
-        for idx, msg in enumerate(smart_split(texts)):
-            if idx == 0:
-                sent_messages.append(await client.send_message(target_chat, msg, reply_parameters=reply_parameters))
-            else:
-                sent_messages.append(await client.send_message(target_chat, msg, reply_parameters=ReplyParameters()))
-        return sent_messages
-    # both texts and media
-    if texts and media:
-        logger.trace(f"Sending {len(media)} media + {len(texts)} texts")
-        # short text, single media
-        if len(texts) < CAPTION_LENGTH and len(media) == 1:
-            if media[0].get("photo"):
-                sent_messages.append(await client.send_photo(target_chat, photo=media[0]["photo"], caption=texts, reply_parameters=reply_parameters))
-            elif media[0].get("video"):
-                sent_messages.append(
-                    await client.send_video(
-                        chat_id=target_chat,
-                        caption=texts,
-                        reply_parameters=reply_parameters,
-                        progress=telegram_uploading,
-                        progress_args=(kwargs.get("progress", False), media[0]["video"], kwargs.get("detail_progress", True)),
-                        **media[0],
-                    )
-                )
-        # long text, single media
-        elif len(texts) >= CAPTION_LENGTH and len(media) == 1:
-            caption_text = smart_split(texts, CAPTION_LENGTH)[0]
-            if media[0].get("photo"):
-                sent_messages.append(await client.send_photo(target_chat, photo=media[0]["photo"], caption=caption_text, reply_parameters=reply_parameters))
-            elif media[0].get("video"):
-                sent_messages.append(
-                    await client.send_video(
-                        chat_id=target_chat,
-                        caption=caption_text,
-                        reply_parameters=reply_parameters,
-                        progress=telegram_uploading,
-                        progress_args=(kwargs.get("progress", False), media[0]["video"], kwargs.get("detail_progress", True)),
-                        **media[0],
-                    )
-                )
-            remaining_texts = texts.removeprefix(caption_text)
-            sent_messages.extend(await send2tg(client, message, target_chat, reply_msg_id=-1, texts=remaining_texts, **kwargs))
-
-        # short text, multiple media [1, 10]
-        elif len(texts) < CAPTION_LENGTH and 1 < len(media) <= 10:
-            group = warp_media_group(media, caption=texts)
-            sent_messages.extend(await client.send_media_group(target_chat, media=group, reply_parameters=reply_parameters))
-
-        # short text, multiple media (10, inf)
-        elif len(texts) < CAPTION_LENGTH and len(media) > 10:
-            media_chunks = [media[i : i + 10] for i in range(0, len(media), 10)]
-            num_chunk = len(media_chunks)
-            # send pure media first, and append captions at the last chunk
-            for idx, batch in enumerate(media_chunks):
-                if idx == 0:  # first chunk
-                    group = warp_media_group(batch)
-                    sent_messages.extend(await client.send_media_group(target_chat, media=group, reply_parameters=reply_parameters))
-                elif idx != num_chunk - 1:  # disbale reply if not the last chunk
-                    group = warp_media_group(batch)
-                    sent_messages.extend(await client.send_media_group(target_chat, media=group, reply_parameters=ReplyParameters()))
-                else:  # last chunk (media <= 10, texts < 1024)
-                    sent_messages.extend(await send2tg(client, message, target_chat, reply_msg_id=-1, texts=texts, media=batch, **kwargs))
-                await asyncio.sleep(cooldown)
-
-        # long text, multiple media [1, 10]
-        elif len(texts) >= CAPTION_LENGTH and 1 < len(media) <= 10:
-            caption_text = smart_split(texts, CAPTION_LENGTH)[0]
-            remaining_texts = texts.removeprefix(caption_text)
-            group = warp_media_group(media, caption=caption_text)
-            sent_messages.extend(await client.send_media_group(target_chat, media=group, reply_parameters=reply_parameters))
-            sent_messages.extend(await send2tg(client, message, target_chat, reply_msg_id=-1, texts=remaining_texts, **kwargs))
-
-        # long text, multiple media (10, inf)
-        else:
-            media_chunks = [media[i : i + 10] for i in range(0, len(media), 10)]
-            num_chunk = len(media_chunks)
-            # send pure media first, and append captions at the last chunk
-            for idx, batch in enumerate(media_chunks):
-                if idx == 0:  # first chunk
-                    group = warp_media_group(batch)
-                    sent_messages.extend(await client.send_media_group(target_chat, media=group, reply_parameters=reply_parameters))
-                elif idx != num_chunk - 1:  # disbale reply if not the last chunk
-                    group = warp_media_group(batch)
-                    sent_messages.extend(await client.send_media_group(target_chat, media=group, reply_parameters=ReplyParameters()))
-                else:  # last chunk (media <= 10, texts >= 1024)
-                    sent_messages.extend(await send2tg(client, message, target_chat, reply_msg_id=-1, texts=texts, media=batch, **kwargs))
-                await asyncio.sleep(cooldown)
-
+    sent_messages: list[Message | None] = []  # return sent messages
+    logger.trace(f"Sending {len(media)} media with {len(texts)} texts")
+    if len(media) == 0:
+        return await send_texts(client, target_chat, reply_parameters, texts=texts, cooldown=cooldown)
+    if len(media) == 1:
+        return await send_single_media(client, target_chat, reply_parameters, media=media[0], texts=texts, cooldown=cooldown, **kwargs)
+
+    caption = smart_split(texts, CAPTION_LENGTH)[0]
+    remaining_texts = texts.removeprefix(caption)
+    if 1 < len(media) <= 10:
+        group = warp_media_group(media, caption=caption)
+        sent_messages.extend(await client.send_media_group(target_chat, media=group, reply_parameters=reply_parameters))
+    else:  # media > 10
+        media_chunks = [media[i : i + 10] for i in range(0, len(media), 10)]
+        num_chunk = len(media_chunks)
+        # send pure media first, and append captions at the last chunk
+        for idx, batch in enumerate(media_chunks):
+            if idx == 0:  # first chunk
+                group = warp_media_group(batch)
+                sent_messages.extend(await client.send_media_group(target_chat, media=group, reply_parameters=reply_parameters))
+            elif idx != num_chunk - 1:  # disbale reply if not the last chunk
+                group = warp_media_group(batch)
+                sent_messages.extend(await client.send_media_group(target_chat, media=group, reply_parameters=ReplyParameters()))
+            else:  # last chunk:  media <= 10, add caption here
+                sent_messages.extend(await send2tg(client, message, target_chat, reply_msg_id=-1, texts=caption, media=batch, cooldown=cooldown, **kwargs))
+            await asyncio.sleep(cooldown)
+    if remaining_texts:
+        sent_messages.extend(await send_texts(client, target_chat, ReplyParameters(), texts=remaining_texts, cooldown=cooldown))
     # clean up
-    logger.trace("Cleaning up media files")
     for x in media:
         for key in ["path", "thumb", "audio", "photo", "video"]:
             if x.get(key) and Path(x[key]).is_file():