Skip to content

Commit

Permalink
refector metadata
Browse files Browse the repository at this point in the history
  • Loading branch information
@5hojib
5hojib committed Jan 1, 2025
1 parent 4974f23 commit 5802a21
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 295 deletions.
37 changes: 12 additions & 25 deletions bot/helper/aeon_utils/metadata_editor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -55,7 +54,7 @@ async def change_metadata(file, key):
"-progress",
"pipe:1",
"-i",
file,
file_path,
"-map_metadata",
"-1",
"-c",
Expand Down Expand Up @@ -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):
Expand Down
128 changes: 13 additions & 115 deletions bot/helper/common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -51,7 +51,6 @@
get_document_type,
is_mkv,
run_ffmpeg_cmd,
run_metadata_cmd,
split_file,
take_ss,
)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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,
Expand All @@ -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:
Expand All @@ -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
Loading

0 comments on commit 5802a21

Please sign in to comment.