Files
moma/renamer/extractors/mediainfo_extractor.py
T
sha 1d6eb9593e feat: Refactor formatting and extraction logic
- Added `langcodes` dependency for improved language handling.
- Replaced `ColorFormatter` with `TextFormatter` for consistent text styling across the application.
- Introduced `TrackFormatter` for better track information formatting.
- Updated `MediaFormatter` to utilize new formatting methods and improved data handling.
- Refactored `MediaExtractor` to enhance data extraction logic and improve readability.
- Removed deprecated `ColorFormatter` methods and replaced them with `TextFormatter` equivalents.
- Added new methods for extracting and formatting audio and subtitle tracks.
- Updated tests to reflect changes in the extraction logic and formatting.
2025-12-26 11:33:24 +00:00

133 lines
5.1 KiB
Python

from pathlib import Path
from pymediainfo import MediaInfo
from collections import Counter
from ..constants import FRAME_CLASSES
import langcodes
class MediaInfoExtractor:
"""Class to extract information from MediaInfo"""
def __init__(self, file_path: Path):
self.file_path = file_path
try:
self.media_info = MediaInfo.parse(file_path)
self.video_tracks = [t for t in self.media_info.tracks if t.track_type == 'Video']
self.audio_tracks = [t for t in self.media_info.tracks if t.track_type == 'Audio']
self.sub_tracks = [t for t in self.media_info.tracks if t.track_type == 'Text']
except Exception:
self.media_info = None
self.video_tracks = []
self.audio_tracks = []
self.sub_tracks = []
def _get_frame_class_from_height(self, height: int) -> str:
"""Get frame class from video height using FRAME_CLASSES constant"""
for frame_class, info in FRAME_CLASSES.items():
if height == info['nominal_height']:
return frame_class
return 'Unclassified'
def extract_frame_class(self) -> str | None:
"""Extract frame class from media info (480p, 720p, 1080p, etc.)"""
if not self.video_tracks:
return 'Unclassified'
height = getattr(self.video_tracks[0], 'height', None)
if height:
return self._get_frame_class_from_height(height)
return 'Unclassified'
def extract_resolution(self) -> tuple[int, int] | None:
"""Extract actual video resolution as (width, height) tuple from media info"""
if not self.video_tracks:
return None
width = getattr(self.video_tracks[0], 'width', None)
height = getattr(self.video_tracks[0], 'height', None)
if width and height:
return width, height
return None
def extract_aspect_ratio(self) -> str | None:
"""Extract video aspect ratio from media info"""
if not self.video_tracks:
return None
aspect_ratio = getattr(self.video_tracks[0], 'display_aspect_ratio', None)
if aspect_ratio:
return str(aspect_ratio)
return None
def extract_hdr(self) -> str | None:
"""Extract HDR info from media info"""
if not self.video_tracks:
return None
profile = getattr(self.video_tracks[0], 'format_profile', '')
if 'HDR' in profile.upper():
return 'HDR'
return None
def extract_audio_langs(self) -> str:
"""Extract audio languages from media info"""
if not self.audio_tracks:
return ''
langs = []
for a in self.audio_tracks:
lang_code = getattr(a, 'language', 'und').lower()
try:
# Try to get the 3-letter code
lang_obj = langcodes.Language.get(lang_code)
alpha3 = lang_obj.to_alpha3()
langs.append(alpha3)
except:
# If conversion fails, use the original code
langs.append(lang_code[:3])
lang_counts = Counter(langs)
audio_langs = [f"{count}{lang}" if count > 1 else lang for lang, count in lang_counts.items()]
return ','.join(audio_langs)
def extract_video_tracks(self) -> list[dict]:
"""Extract video track data"""
tracks = []
for v in self.video_tracks[:2]: # Up to 2 videos
track_data = {
'codec': getattr(v, 'format', None) or getattr(v, 'codec', None) or 'unknown',
'width': getattr(v, 'width', None),
'height': getattr(v, 'height', None),
'bitrate': getattr(v, 'bit_rate', None),
'fps': getattr(v, 'frame_rate', None),
'profile': getattr(v, 'format_profile', None),
}
tracks.append(track_data)
return tracks
def extract_audio_tracks(self) -> list[dict]:
"""Extract audio track data"""
tracks = []
for a in self.audio_tracks[:3]: # Up to 3 audios
track_data = {
'codec': getattr(a, 'format', None) or getattr(a, 'codec', None) or 'unknown',
'channels': getattr(a, 'channel_s', None),
'language': getattr(a, 'language', None) or 'und',
'bitrate': getattr(a, 'bit_rate', None),
}
tracks.append(track_data)
return tracks
def extract_subtitle_tracks(self) -> list[dict]:
"""Extract subtitle track data"""
tracks = []
for s in self.sub_tracks[:3]: # Up to 3 subs
track_data = {
'language': getattr(s, 'language', None) or 'und',
'format': getattr(s, 'format', None) or getattr(s, 'codec', None) or 'unknown',
}
tracks.append(track_data)
return tracks
def extract_tracks(self) -> dict:
"""Extract media track information as data"""
return {
'video_tracks': self.extract_video_tracks(),
'audio_tracks': self.extract_audio_tracks(),
'subtitle_tracks': self.extract_subtitle_tracks(),
}