Skip to content

Commit

Permalink
fix: connect lost (#217)
Browse files Browse the repository at this point in the history
feat: support forward special topic
fix: forward private text
fix: support topic_id filter
feat: support forward comment
feat: support caption entities
fix: connect lost
  • Loading branch information
tangyoha authored Jan 6, 2025
1 parent 0d82a96 commit b5f9504
Show file tree
Hide file tree
Showing 15 changed files with 567 additions and 168 deletions.
34 changes: 23 additions & 11 deletions media_downloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,14 @@
report_bot_download_status,
set_max_concurrent_transmissions,
set_meta_data,
update_cloud_upload_stat,
upload_telegram_chat,
)
from module.web import init_web
from utils.format import truncate_filename, validate_title
from utils.log import LogFilter
from utils.meta import print_meta
from utils.meta_data import MetaData
from utils.updates import check_for_updates

logging.basicConfig(
level=logging.INFO,
Expand Down Expand Up @@ -234,6 +234,9 @@ async def _get_media_meta(
if caption:
caption = validate_title(caption)
app.set_caption_name(chat_id, message.media_group_id, caption)
app.set_caption_entities(
chat_id, message.media_group_id, message.caption_entities
)
else:
caption = app.get_caption_name(chat_id, message.media_group_id)

Expand Down Expand Up @@ -326,7 +329,12 @@ async def download_task(
not node.upload_telegram_chat_id
and download_status is DownloadStatus.SuccessDownload
):
if await app.upload_file(file_name):
ui_file_name = file_name
if app.hide_file_name:
ui_file_name = f"****{os.path.splitext(file_name)[-1]}"
if await app.upload_file(
file_name, update_cloud_upload_stat, (node, message.id, ui_file_name)
):
node.upload_success_count += 1

await report_bot_download_status(
Expand Down Expand Up @@ -563,6 +571,9 @@ async def download_chat_task(
if caption:
caption = validate_title(caption)
app.set_caption_name(node.chat_id, message.media_group_id, caption)
app.set_caption_entities(
node.chat_id, message.media_group_id, message.caption_entities
)
else:
caption = app.get_caption_name(node.chat_id, message.media_group_id)
set_meta_data(meta_data, message, caption)
Expand All @@ -574,14 +585,15 @@ async def download_chat_task(
await add_download_task(message, node)
else:
node.download_status[message.id] = DownloadStatus.SkipDownload
await upload_telegram_chat(
client,
node.upload_user,
app,
node,
message,
DownloadStatus.SkipDownload,
)
if message.media_group_id:
await upload_telegram_chat(
client,
node.upload_user,
app,
node,
message,
DownloadStatus.SkipDownload,
)

chat_download_config.need_check = True
chat_download_config.total_task = node.total_task
Expand Down Expand Up @@ -676,7 +688,7 @@ def main():
for task in tasks:
task.cancel()
logger.info(_t("Stopped!"))
check_for_updates(app.proxy)
# check_for_updates(app.proxy)
logger.info(f"{_t('update config')}......")
app.update_config()
logger.success(
Expand Down
1 change: 1 addition & 0 deletions media_downloader.spec
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ exe = EXE(
target_arch=None,
codesign_identity=None,
entitlements_file=None,
contents_directory='.',
)
coll = COLLECT(
exe,
Expand Down
63 changes: 59 additions & 4 deletions module/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import List, Optional, Union
from typing import Callable, List, Optional, Union

from loguru import logger
from ruamel import yaml
Expand Down Expand Up @@ -76,10 +76,21 @@ class UploadProgressStat:
upload_size: int
start_time: float
last_stat_time: float
each_second_total_upload: int
upload_speed: float


@dataclass
class CloudDriveUploadStat:
"""Cloud drive upload task"""

file_name: str
transferred: str
total: str
percentage: str
speed: str
eta: str


class QueryHandlerStr:
"""Query handler"""

Expand Down Expand Up @@ -122,6 +133,7 @@ def __init__(
bot=None,
task_type: TaskType = TaskType.Download,
task_id: int = 0,
topic_id: int = 0,
):
self.chat_id = chat_id
self.from_user_id = from_user_id
Expand Down Expand Up @@ -158,6 +170,9 @@ def __init__(
self.download_status: dict = {}
self.upload_status: dict = {}
self.upload_stat_dict: dict = {}
self.topic_id = topic_id
self.reply_to_message = None
self.cloud_drive_upload_stat_dict: dict = {}

def skip_msg_id(self, msg_id: int):
"""Skip if message id out of range"""
Expand Down Expand Up @@ -376,6 +391,7 @@ def __init__(
self.cloud_drive_config = CloudDriveConfig()
self.hide_file_name = False
self.caption_name_dict: dict = {}
self.caption_entities_dict: dict = {}
self.max_concurrent_transmissions: int = 1
self.web_host: str = "0.0.0.0"
self.web_port: int = 5000
Expand All @@ -394,6 +410,7 @@ def __init__(
self.enable_download_txt: bool = False

self.forward_limit_call = LimitCall(max_limit_call_times=33)

self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.loop)

Expand Down Expand Up @@ -643,7 +660,12 @@ def assign_app_data(self, app_data: dict) -> bool:
] = True
return True

async def upload_file(self, local_file_path: str) -> bool:
async def upload_file(
self,
local_file_path: str,
progress_callback: Callable = None,
progress_args: tuple = (),
) -> bool:
"""Upload file"""

if not self.cloud_drive_config.enable_upload_file:
Expand All @@ -652,7 +674,11 @@ async def upload_file(self, local_file_path: str) -> bool:
ret: bool = False
if self.cloud_drive_config.upload_adapter == "rclone":
ret = await CloudDrive.rclone_upload_file(
self.cloud_drive_config, self.save_path, local_file_path
self.cloud_drive_config,
self.save_path,
local_file_path,
progress_callback,
progress_args,
)
elif self.cloud_drive_config.upload_adapter == "aligo":
ret = await self.loop.run_in_executor(
Expand Down Expand Up @@ -924,6 +950,35 @@ def get_caption_name(

return str(self.caption_name_dict[chat_id][media_group_id])

def set_caption_entities(
self, chat_id: Union[int, str], media_group_id: Optional[str], caption_entities
):
"""
set caption entities map
"""
if not media_group_id:
return

if chat_id in self.caption_entities_dict:
self.caption_entities_dict[chat_id][media_group_id] = caption_entities
else:
self.caption_entities_dict[chat_id] = {media_group_id: caption_entities}

def get_caption_entities(
self, chat_id: Union[int, str], media_group_id: Optional[str]
):
"""
get caption entities map
"""
if (
not media_group_id
or chat_id not in self.caption_entities_dict
or media_group_id not in self.caption_entities_dict[chat_id]
):
return None

return self.caption_entities_dict[chat_id][media_group_id]

def set_download_id(
self, node: TaskNode, message_id: int, download_status: DownloadStatus
):
Expand Down
Loading

0 comments on commit b5f9504

Please sign in to comment.