From 5802a213ddc00fa20d785996d26d8a5df769c077 Mon Sep 17 00:00:00 2001 From: 5hojib Date: Wed, 1 Jan 2025 23:11:31 +0600 Subject: [PATCH] refector metadata --- bot/helper/aeon_utils/metadata_editor.py | 37 ++--- bot/helper/common.py | 128 ++-------------- bot/helper/ext_utils/media_utils.py | 183 ++++------------------- 3 files changed, 53 insertions(+), 295 deletions(-) diff --git a/bot/helper/aeon_utils/metadata_editor.py b/bot/helper/aeon_utils/metadata_editor.py index 02c10ead7..60250d694 100644 --- a/bot/helper/aeon_utils/metadata_editor.py +++ b/bot/helper/aeon_utils/metadata_editor.py @@ -33,19 +33,18 @@ async def get_streams(file): return None -async def change_metadata(file, key): - LOGGER.info(f"Starting metadata modification for file: {file}") - temp_file = f"{file}.temp.mkv" - - streams = await get_streams(file) +async def get_metadata_cmd(file_path, key): + """Processes a single file to update metadata.""" + temp_file = f"{file_path}.temp.mkv" + streams = await get_streams(file_path) if not streams: - return + return None, None - languages = {} - for stream in streams: - stream_index = stream["index"] - if "tags" in stream and "language" in stream["tags"]: - languages[stream_index] = stream["tags"]["language"] + languages = { + stream["index"]: stream["tags"]["language"] + for stream in streams + if "tags" in stream and "language" in stream["tags"] + } cmd = [ "xtra", @@ -55,7 +54,7 @@ async def change_metadata(file, key): "-progress", "pipe:1", "-i", - file, + file_path, "-map_metadata", "-1", "-c", @@ -130,19 +129,7 @@ async def change_metadata(file, key): cmd.extend(["-map", f"0:{stream_index}"]) cmd.append(temp_file) - - process = await create_subprocess_exec(*cmd, stderr=PIPE, stdout=PIPE) - stdout, stderr = await process.communicate() - - if process.returncode != 0: - err = stderr.decode().strip() - LOGGER.error(err) - LOGGER.error(f"Error modifying metadata for file: {file}") - return - - os.replace(temp_file, file) - LOGGER.info(f"Metadata modified successfully for file: {file}") - return + return cmd, temp_file async def add_attachment(file, attachment_path): diff --git a/bot/helper/common.py b/bot/helper/common.py index 8d16ad1a7..fad2008d1 100644 --- a/bot/helper/common.py +++ b/bot/helper/common.py @@ -24,7 +24,7 @@ ) from bot.core.aeon_client import TgClient from bot.core.config_manager import Config -from bot.helper.aeon_utils.metadata_editor import get_streams +from bot.helper.aeon_utils.metadata_editor import get_streams, get_metadata_cmd from .ext_utils.bot_utils import get_size_bytes, new_task, sync_to_async from .ext_utils.bulk_links import extract_bulk_links @@ -51,7 +51,6 @@ get_document_type, is_mkv, run_ffmpeg_cmd, - run_metadata_cmd, split_file, take_ss, ) @@ -1269,7 +1268,7 @@ async def proceed_ffmpeg(self, dl_path, gid): LOGGER.info(f"Running ffmpeg cmd for: {file_path}") cmd[index + 1] = file_path self.subsize = self.size - res = await run_ffmpeg_cmd(self, cmd, file_path) + res = await run_ffmpeg_cmd(self, cmd, file_path, True) self.subproc = None if res and delete_files: await remove(file_path) @@ -1312,7 +1311,7 @@ async def proceed_ffmpeg(self, dl_path, gid): LOGGER.info(f"Running ffmpeg cmd for: {f_path}") self.subsize = await aiopath.getsize(f_path) self.subname = file_ - res = await run_ffmpeg_cmd(self, cmd, f_path) + res = await run_ffmpeg_cmd(self, cmd, f_path, True) self.suproc = None if res and delete_files: await remove(f_path) @@ -1326,112 +1325,12 @@ async def proceed_ffmpeg(self, dl_path, gid): cpu_eater_lock.release() return dl_path - async def proceed_metadata(self, up_dir, gid): + async def proceed_metadata(self, dl_path, gid): key = self.metadata async with task_dict_lock: task_dict[self.mid] = FFmpegStatus(self, gid, "Metadata") checked = False - - async def process_file(file_path, key): - """Processes a single file to update metadata.""" - temp_file = f"{file_path}.temp.mkv" - streams = await get_streams(file_path) - if not streams: - return "" - - languages = { - stream["index"]: stream["tags"]["language"] - for stream in streams - if "tags" in stream and "language" in stream["tags"] - } - - cmd = [ - "xtra", - "-hide_banner", - "-loglevel", - "error", - "-progress", - "pipe:1", - "-i", - file_path, - "-map_metadata", - "-1", - "-c", - "copy", - "-metadata:s:v:0", - f"title={key}", - "-metadata", - f"title={key}", - ] - - audio_index = 0 - subtitle_index = 0 - first_video = False - - for stream in streams: - stream_index = stream["index"] - stream_type = stream["codec_type"] - - if stream_type == "video": - if not first_video: - cmd.extend(["-map", f"0:{stream_index}"]) - first_video = True - cmd.extend([f"-metadata:s:v:{stream_index}", f"title={key}"]) - if stream_index in languages: - cmd.extend( - [ - f"-metadata:s:v:{stream_index}", - f"language={languages[stream_index]}", - ], - ) - elif stream_type == "audio": - cmd.extend( - [ - "-map", - f"0:{stream_index}", - f"-metadata:s:a:{audio_index}", - f"title={key}", - ], - ) - if stream_index in languages: - cmd.extend( - [ - f"-metadata:s:a:{audio_index}", - f"language={languages[stream_index]}", - ], - ) - audio_index += 1 - elif stream_type == "subtitle": - codec_name = stream.get("codec_name", "unknown") - if codec_name in ["webvtt", "unknown"]: - LOGGER.warning( - f"Skipping unsupported subtitle metadata modification: {codec_name} for stream {stream_index}", - ) - else: - cmd.extend( - [ - "-map", - f"0:{stream_index}", - f"-metadata:s:s:{subtitle_index}", - f"title={key}", - ], - ) - if stream_index in languages: - cmd.extend( - [ - f"-metadata:s:s:{subtitle_index}", - f"language={languages[stream_index]}", - ], - ) - subtitle_index += 1 - else: - cmd.extend(["-map", f"0:{stream_index}"]) - - cmd.append(temp_file) - return cmd, temp_file - - async def process_directory(directory, key): - """Processes all MKV files in a directory.""" + async def _process_directory(directory, key): for dirpath, _, files in await sync_to_async( walk, directory, @@ -1440,7 +1339,7 @@ async def process_directory(directory, key): for file_ in files: file_path = ospath.join(dirpath, file_) if is_mkv(file_path): - cmd, temp_file = await process_file(file_path, key) + cmd, temp_file = await get_metadata_cmd(file_path, key) if not cmd: return "" if not checked: @@ -1450,24 +1349,23 @@ async def process_directory(directory, key): return "" self.subsize = await aiopath.getsize(file_path) self.subname = file_ - await run_metadata_cmd(self, cmd) + await run_ffmpeg_cmd(self, cmd) self.subproc = None os.replace(temp_file, file_path) return None - if self.is_file: - if is_mkv(up_dir): - cmd, temp_file = await process_file(up_dir, key) + if is_mkv(dl_path): + cmd, temp_file = await get_metadata_cmd(dl_path, key) if cmd: checked = True await cpu_eater_lock.acquire() self.subsize = self.size - await run_metadata_cmd(self, cmd) + await run_ffmpeg_cmd(self, cmd) self.subproc = None - os.replace(temp_file, up_dir) + os.replace(temp_file, dl_path) else: - await process_directory(up_dir, key) + await _process_directory(dl_path, key) if checked: cpu_eater_lock.release() - return up_dir + return dl_path \ No newline at end of file diff --git a/bot/helper/ext_utils/media_utils.py b/bot/helper/ext_utils/media_utils.py index ddacafa9e..1d0731aa1 100644 --- a/bot/helper/ext_utils/media_utils.py +++ b/bot/helper/ext_utils/media_utils.py @@ -504,7 +504,6 @@ async def split_file( async with listener.subprocess_lock: if listener.is_cancelled: return False - code = listener.subproc.returncode if code == -9: listener.is_cancelled = True return False @@ -677,117 +676,25 @@ async def create_sample_video(listener, video_file, sample_duration, part_durati await remove(output_file) return False - """finished_segments = [] - await makedirs(f"{dir}/segments/", exist_ok=True) - ext = name.rsplit(".", 1)[-1] - for index, (start_time, end_time) in enumerate(segments, start=1): - output_seg = f"{dir}/segments/segment{index}.{ext}" - cmd = [ - "xtra", - "-hide_banner", - "-loglevel", - "error", - "-i", - video_file, - "-ss", - f"{start_time}", - "-to", - f"{end_time}", - "-c", - "copy", - output_seg, - ] - if listener.is_cancelled: - return False - listener.subproc = await create_subprocess_exec(*cmd, stderr=PIPE) - _, stderr = await listener.subproc.communicate() - if listener.is_cancelled: - return False - code = listener.subproc.returncode - if code == -9: - listener.is_cancelled = True - return False - elif code != 0: - try: - stderr = stderr.decode().strip() - except Exception: - stderr = "Unable to decode the error!" - LOGGER.error( - f"{stderr}. Something went wrong while splitting file for sample video, mostly file is corrupted. Path: {video_file}" - ) - if await aiopath.exists(output_file): - await remove(output_file) - return False +async def run_ffmpeg_cmd(listener, cmd, path=None, is_ffmpeg=False): + if is_ffmpeg and path: + base_name, ext = ospath.splitext(path) + dir, base_name = base_name.rsplit("/", 1) + output_file = cmd[-1] + if output_file != "mltb" and output_file.startswith("mltb"): + oext = ospath.splitext(output_file)[-1] + if ext == oext: + base_name = f"ffmpeg.{base_name}" + else: + ext = oext else: - finished_segments.append(f"file '{output_seg}'") - - segments_file = f"{dir}/segments.txt" - - async with aiopen(segments_file, "w+") as f: - await f.write("\n".join(finished_segments)) - - cmd = [ - "xtra", - "-hide_banner", - "-loglevel", - "error", - "-f", - "concat", - "-safe", - "0", - "-i", - segments_file, - "-c:v", - "libx264", - "-c:a", - "aac", - "-threads", - f"{cpu_count() // 2}", - output_file, - ] - if listener.is_cancelled: - return False - listener.subproc = await create_subprocess_exec(*cmd, stderr=PIPE) - _, stderr = await listener.subproc.communicate() - if listener.is_cancelled: - return False - code = listener.subproc.returncode - if code == -9: - listener.is_cancelled = True - return False - elif code != 0: - try: - stderr = stderr.decode().strip() - except Exception: - stderr = "Unable to decode the error!" - LOGGER.error( - f"{stderr}. Something went wrong while creating sample video, mostly file is corrupted. Path: {video_file}" - ) - if await aiopath.exists(output_file): - await remove(output_file) - await gather(remove(segments_file), rmtree(f"{dir}/segments")) - return False - await gather(remove(segments_file), rmtree(f"{dir}/segments")) - return output_file""" - return None - - -async def run_ffmpeg_cmd(listener, cmd, path): - base_name, ext = ospath.splitext(path) - dir, base_name = base_name.rsplit("/", 1) - output_file = cmd[-1] - if output_file != "mltb" and output_file.startswith("mltb"): - oext = ospath.splitext(output_file)[-1] - if ext == oext: base_name = f"ffmpeg.{base_name}" - else: - ext = oext - else: - base_name = f"ffmpeg.{base_name}" - output = f"{dir}/{base_name}{ext}" - cmd[-1] = output + output = f"{dir}/{base_name}{ext}" + cmd[-1] = output + if listener.is_cancelled: return False + async with listener.subprocess_lock: listener.subproc = await create_subprocess_exec( *cmd, @@ -795,66 +702,32 @@ async def run_ffmpeg_cmd(listener, cmd, path): stderr=PIPE, ) code = await listener.subproc.wait() + async with listener.subprocess_lock: if listener.is_cancelled: return False + if code == 0: - return output + return output if is_ffmpeg and path else True if code == -9: listener.is_cancelled = True return False + try: stderr = (await listener.subproc.stderr.read()).decode().strip() except Exception: stderr = "Unable to decode the error!" - LOGGER.error( - f"{stderr}. Something went wrong while running ffmpeg cmd, mostly file requires different/specific arguments. Path: {path}", - ) - if await aiopath.exists(output): - await remove(output) - return False - -async def run_metadata_cmd(listener, cmd): - # base_name, ext = ospath.splitext(path) - # dir, base_name = base_name.rsplit("/", 1) - # output_file = cmd[-1] - # if output_file != "mltb" and output_file.startswith("mltb"): - # oext = ospath.splitext(output_file)[-1] - # if ext == oext: - # base_name = f"ffmpeg.{base_name}" - # else: - # ext = oext - # else: - # base_name = f"ffmpeg.{base_name}" - # output = f"{dir}/{base_name}{ext}" - # cmd[-1] = output - if listener.is_cancelled: - return False - async with listener.subprocess_lock: - listener.subproc = await create_subprocess_exec( - *cmd, - stdout=PIPE, - stderr=PIPE, + if is_ffmpeg and path: + LOGGER.error( + f"{stderr}. Something went wrong while running ffmpeg cmd, mostly file requires different/specific arguments. Path: {path}", + ) + if await aiopath.exists(output): + await remove(output) + else: + LOGGER.error( + f"{stderr}. Something went wrong while running metadata cmd, mostly file requires different/specific arguments.", ) - code = await listener.subproc.wait() - async with listener.subprocess_lock: - if listener.is_cancelled: - return False - if code == 0: - return True - if code == -9: - listener.is_cancelled = True - return False - try: - stderr = (await listener.subproc.stderr.read()).decode().strip() - except Exception: - stderr = "Unable to decode the error!" - LOGGER.error( - f"{stderr}. Something went wrong while running metadata cmd, mostly file requires different/specific arguments.", - ) - # if await aiopath.exists(output): - # await remove(output) return False