From bf78b32c6b9debca191dba09237e578f7b53d8f2 Mon Sep 17 00:00:00 2001
From: tangyoha <tangyoha@outlook.com>
Date: Thu, 30 Nov 2023 19:41:34 +0800
Subject: [PATCH] feat: support caption replace

---
 .github/workflows/docker-publish.yml |  5 ++---
 module/app.py                        | 16 +++++++++++++++
 module/pyrogram_extension.py         | 29 +++++++++++++++++++++++++---
 3 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml
index 6ee24e05..2633f901 100644
--- a/.github/workflows/docker-publish.yml
+++ b/.github/workflows/docker-publish.yml
@@ -4,6 +4,7 @@ on:
   push:
     branches:
       - master
+      - plus
     tags:
       - 'v*'
 
@@ -75,6 +76,4 @@ jobs:
         push: true
         platforms: linux/amd64,linux/arm64,linux/arm/v7,linux/arm/v6,linux/386,linux/ppc64le
         target: runtime-image
-        tags: |
-          ${{ secrets.DOCKER_HUB_USERNAME }}/telegram_media_downloader:latest
-          ${{ secrets.DOCKER_HUB_USERNAME }}/telegram_media_downloader:${{ github.ref_name }}
+        tags: ${{ secrets.DOCKER_HUB_USERNAME }}/telegram_media_downloader:${{ github.ref_name }}
diff --git a/module/app.py b/module/app.py
index 28c340e2..0990674f 100644
--- a/module/app.py
+++ b/module/app.py
@@ -408,6 +408,8 @@ def __init__(
 
         self.caption_replace_dict: yaml.comments.CommentedMap = {}
         self.default_forward_caption = None
+        self.caption_regex_replace_dict: yaml.comments.CommentedMap = {}
+        self.default_forward_additional_caption = None
         self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(self.loop)
 
@@ -566,6 +568,20 @@ def assign_config(self, _config: dict) -> bool:
             _config, "default_forward_caption", self.default_forward_caption, str
         )
 
+        self.caption_regex_replace_dict = get_config(
+            _config,
+            "caption_replace_regex",
+            self.caption_regex_replace_dict,
+            yaml.comments.CommentedMap,
+        )
+
+        self.default_forward_additional_caption = get_config(
+            _config,
+            "default_forward_additional_caption",
+            self.default_forward_additional_caption,
+            str,
+        )
+
         if _config.get("chat"):
             chat = _config["chat"]
             for item in chat:
diff --git a/module/pyrogram_extension.py b/module/pyrogram_extension.py
index 41e2bc43..f3d822f5 100644
--- a/module/pyrogram_extension.py
+++ b/module/pyrogram_extension.py
@@ -2,6 +2,7 @@
 
 import asyncio
 import os
+import re
 import secrets
 import struct
 import time
@@ -107,7 +108,11 @@ def get_media_obj(
 
 
 def replace_caption(
-    caption: Optional[str], caption_replace_dict, default_caption: Optional[str] = None
+    caption: Optional[str],
+    caption_replace_dict,
+    default_caption: Optional[str] = None,
+    caption_regex_replace_dict=None,
+    default_additional_caption: Optional[str] = None,
 ):
     """
     Replaces certain items in a caption string
@@ -125,6 +130,16 @@ def replace_caption(
             caption = caption.replace(item, caption_replace_dict[item])
     else:
         caption = default_caption
+
+    if not caption:
+        return default_additional_caption
+
+    if caption_regex_replace_dict:
+        for item in caption_regex_replace_dict:
+            caption = re.sub(item, caption_regex_replace_dict[item], caption)
+
+    if default_additional_caption:
+        caption += default_additional_caption
     return caption
 
 
@@ -368,7 +383,11 @@ async def _upload_signal_message(
         )
 
     caption = replace_caption(
-        message.caption, app.caption_replace_dict, app.default_forward_caption
+        message.caption,
+        app.caption_replace_dict,
+        app.default_forward_caption,
+        app.caption_regex_replace_dict,
+        app.default_forward_additional_caption,
     )
     if message.video:
         # Download thumbnail
@@ -499,7 +518,11 @@ async def forward_multi_media(
         caption = app.get_caption_name(node.chat_id, message.media_group_id)
 
     caption = replace_caption(
-        caption, app.caption_replace_dict, app.default_forward_caption
+        caption,
+        app.caption_replace_dict,
+        app.default_forward_caption,
+        app.caption_regex_replace_dict,
+        app.default_forward_additional_caption,
     )
 
     media_obj = get_media_obj(message, file_name, caption)