Rename SRT files to match video files for TV series

rename_srt.py

import os
import sys
import unittest
import re
from pathlib import Path
import json


class Constants(object):
    def __init__(self):
        self.ALLOWED_VIDEO_EXTENSIONS = ["AVI", "FLV", "WMV", "MP4", "MOV", "MKV"]
        self.ALLOWED_SRT_EXTENSIONS = ["SRT"]


class RenameSRTFiles(Constants):
    def __init__(self):
        super().__init__()
        self.all_files = self.get_all_files()
        
    def get_all_files(self):
        return os.listdir()

    def get_file_extension(self, filename):
        extension = Path(filename).suffix.strip()
        if extension != "":
            return extension[1:].upper()
        return None

    def get_filename_without_extension(self, filename):
        return Path(filename).stem

    def get_all_files_of_specific_extensions(self, allowed_extensions):
        files = []
        for filename in self.all_files:
            extension = self.get_file_extension(filename)
            if extension and extension in allowed_extensions:
                files.append(filename)
        return files
    

    def get_all_videos(self):
        return self.get_all_files_of_specific_extensions(self.ALLOWED_VIDEO_EXTENSIONS)
    
    def get_all_srt_files(self):
        return self.get_all_files_of_specific_extensions(self.ALLOWED_SRT_EXTENSIONS)

    def get_episode_from_filename(self, filename):
        episode_found = re.search(r'e\d+', filename)
        if episode_found:
            return episode_found.group(0)
        return None

    def get_same_episode_file_group(self):
        srts = self.get_all_srt_files()
        videos = self.get_all_videos()
        data = []
        for video in videos:
            video_episode = self.get_episode_from_filename(video)
            if video_episode:
                for srt in srts:
                    srt_episode = self.get_episode_from_filename(srt)
                    if video_episode == srt_episode:
                        data.append([video, srt])
                        break
        return data

    def perform_rename_operation(self):
        try:
            group_files = self.get_same_episode_file_group()
            for pair in group_files:
                video_file, srt_file = pair
                srt_name = self.get_filename_without_extension(video_file)+"."+self.ALLOWED_SRT_EXTENSIONS[0]
                os.rename(srt_file, srt_name)
        except:
            return False
        return True
                    


class TestRenameSRTFiles(unittest.TestCase):
    def setUp(self):
        self.rename_srt_files = RenameSRTFiles()

    def test_init(self):
        all_files = self.rename_srt_files.all_files
        assert type(all_files) == type([])

    def test_get_all_files_returns_list(self):
        all_files = self.rename_srt_files.get_all_files()        
        assert type(all_files) == type([])

    def test_get_all_videos(self):
        filenames = self.rename_srt_files.get_all_videos()
        for filename in filenames:
            file_extension = filename.rsplit(".", 1)[1]
            assert (file_extension.upper() in self.rename_srt_files.ALLOWED_VIDEO_EXTENSIONS), \
                   file_extension.upper()

    def test_get_all_srt_files(self):
        filenames = self.rename_srt_files.get_all_srt_files()
        for filename in filenames:
            file_extension = filename.rsplit(".", 1)[1]            
            assert (file_extension.upper() in self.rename_srt_files.ALLOWED_SRT_EXTENSIONS), \
                   file_extension.upper()

    def test_get_same_episode_file_group(self):
        group_files = self.rename_srt_files.get_same_episode_file_group()
        assert type(group_files) == type([])

    def test_perform_rename_operation(self):
        rename_operation = self.rename_srt_files.perform_rename_operation()
        assert rename_operation == True


if __name__ == '__main__':
    unittest.main()