feat: Add HDR extraction from filename and update tests

This commit is contained in:
sHa
2025-12-26 13:51:12 +00:00
parent 91df347727
commit 3ecffc0d3c
4 changed files with 31 additions and 7 deletions

View File

@@ -38,7 +38,8 @@ class MediaExtractor:
('MediaInfo', lambda: self.mediainfo_extractor.extract_aspect_ratio()) ('MediaInfo', lambda: self.mediainfo_extractor.extract_aspect_ratio())
], ],
'hdr': [ 'hdr': [
('MediaInfo', lambda: self.mediainfo_extractor.extract_hdr()) ('MediaInfo', lambda: self.mediainfo_extractor.extract_hdr()),
('Filename', lambda: self.filename_extractor.extract_hdr())
], ],
'audio_langs': [ 'audio_langs': [
('MediaInfo', lambda: self.mediainfo_extractor.extract_audio_langs()) ('MediaInfo', lambda: self.mediainfo_extractor.extract_audio_langs())
@@ -46,9 +47,6 @@ class MediaExtractor:
'meta_type': [ 'meta_type': [
('Metadata', lambda: self.metadata_extractor.extract_meta_type()) ('Metadata', lambda: self.metadata_extractor.extract_meta_type())
], ],
'meta_description': [
('Metadata', lambda: self.metadata_extractor.extract_meta_description())
],
'file_size': [ 'file_size': [
('FileInfo', lambda: self.fileinfo_extractor.extract_size()) ('FileInfo', lambda: self.fileinfo_extractor.extract_size())
], ],

View File

@@ -77,3 +77,15 @@ class FilenameExtractor:
return 'Unclassified' return 'Unclassified'
return 'Unclassified' return 'Unclassified'
def extract_hdr(self) -> str | None:
"""Extract HDR information from filename"""
# Check for SDR first - indicates no HDR
if re.search(r'\bSDR\b', self.file_name, re.IGNORECASE):
return None
# Check for HDR, but not NoHDR
if re.search(r'\bHDR\b', self.file_name, re.IGNORECASE) and not re.search(r'\bNoHDR\b', self.file_name, re.IGNORECASE):
return 'HDR'
return None

View File

@@ -208,4 +208,5 @@ The long title (2008) SD 720p UKR.avi
The long title (2008) UHD 1440p ENG.mp4 The long title (2008) UHD 1440p ENG.mp4
The long title (2008) UHD 1440 ENG.mp4 The long title (2008) UHD 1440 ENG.mp4
The long title (2008) 8K 4320p ENG.mp4 The long title (2008) 8K 4320p ENG.mp4
Troll 2 (2025) WEB-DL 2160p HDR Ukr Nor [Hurtom].mkv
Moana 2 (2024) MA WEB-DL 2160p SDR Ukr Eng [Hurtom].mkv

View File

@@ -67,3 +67,16 @@ def test_extract_frame_class(filename):
# Should be one of the valid frame classes or 'Unclassified' # Should be one of the valid frame classes or 'Unclassified'
valid_classes = set(FRAME_CLASSES.keys()) | {'Unclassified'} valid_classes = set(FRAME_CLASSES.keys()) | {'Unclassified'}
assert frame_class in valid_classes assert frame_class in valid_classes
@pytest.mark.parametrize("filename", load_test_filenames())
def test_extract_hdr(filename):
"""Test HDR extraction from filename"""
file_path = Path(filename)
extractor = FilenameExtractor(file_path)
hdr = extractor.extract_hdr()
# Print filename and extracted HDR clearly
print(f"\nFilename: \033[1;36m{filename}\033[0m")
print(f"Extracted HDR: \033[1;32m{hdr}\033[0m")
# HDR should be 'HDR' or None
assert hdr is None or hdr == 'HDR'