1
\$\begingroup\$

This code unzips the downloaded subtitle files and maps them to their corresponding video file. Currently, after pairing video files with their corresponding subtitle files, it links them by renaming the subtitle file to the same file name as that of its corresponding video file.

#!/usr/bin/env python3
import os, os.path
import shutil
from glob import glob
from PTN import parse
from zipfile import ZipFile


class Subtitles:
    """
    Class to perform subtitle-related operations.
    """
    def __init__(self, subtitleDir, targetDir, tempStorageDir):
        """
        Function to get file paths.

        subtitleDir:    Path to source directory for subtitles.
        targetDir: Path to directory which contains files
                        to which subtitles need to be mapped.
        tempStorageDir: Path to directory to be used as storage
                        for saving files temporarily.
        """
        self.subtitleDir = subtitleDir
        self.targetDir = targetDir
        self.tempStorageDir = tempStorageDir
        self._videoFileFormats = ['3gp', 'avi', 'mkv', 'mp4', 'webm']

    def getTargetFileNames(self):
        """
        Function to get a list a file names which are to be
        mapped to corresponding subtitle file.
        """
        currentWorkingDir = os.getcwd()
        os.chdir(self.targetDir)
        targetFileNames = list()
        for ext in self._videoFileFormats:
            targetFileNames.extend(glob(f'*.{ext}'))
        os.chdir(currentWorkingDir)
        return targetFileNames

    def getZipFileNames(self):
        """
        Function to get a list of zip files present in subtitles
        directory.
        """
        zipFileNames = glob(os.path.join(self.subtitleDir, '*.zip'))
        return zipFileNames

    def unzipSubtitles(self, zipFilePath):
        """
        Function to unzip a file, extract all subtitles from it
        to temporary storage directory and return list of
        subtitle filenames.
        """
        with ZipFile(zipFilePath, 'r') as zipObj:
           listOfFileNames = zipObj.namelist()
           listOfSubtitleFileNames = list()
           for fileName in listOfFileNames:
               if fileName.endswith('.srt'):
                   listOfSubtitleFileNames.append(fileName)
                   zipObj.extract(fileName, self.tempStorageDir)
        return listOfSubtitleFileNames

    def genFileMetadataSummary(self, metadata):
        """
        Function to generate summary of metadata in the format
        <title>.S<season>E<episode>.<quality> and return it.
        """
        fileMetadataSummary = metadata['title'] + '.S' \
              + str(metadata['season']) + 'E' \
              + str(metadata['episode']) + '.' \
              + metadata['quality']
        return fileMetadataSummary

    def cacheFileMetadata(self, listOfFileNames):
        """
        Function to build a dictionary containing metadata
        about each of the files and return it.
        """
        fileMetadata = dict()
        for fileName in listOfFileNames:
            metadata = parse(fileName)
            metadata['filename'] = fileName[:-4]
            fileMetadataSummary = self.genFileMetadataSummary(metadata)
            fileMetadata[fileMetadataSummary] = metadata
        return fileMetadata

    def cleanTempStorageDir(self, listOfFileNames):
        """
        Function to remove a list of files from a temporary
        storage directory.
        """
        for fileName in listOfFileNames:
            try:
                os.remove(os.path.join(self.tempStorageDir, fileName))
            except FileNotFoundError:
                pass

    def mapSubtitlesToFiles(self):
        """
        Function to map a subtitle to corresponding file.
        """
        targetFileNames = self.getTargetFileNames()
        subtitleFileNames = list()

        zipFileNames = self.getZipFileNames()
        for zipFileName in zipFileNames:
            subtitleFileNames.extend(self.unzipSubtitles(zipFileName))

        subtitleMetadata = self.cacheFileMetadata(subtitleFileNames)
        targetMetadata = self.cacheFileMetadata(targetFileNames)

        for targetFileMetadataSummary in targetMetadata:
            try:
                sourceFileName = subtitleMetadata[targetFileMetadataSummary]['filename']+'.srt'
                print(f'Subtitle for \'{targetFileMetadataSummary}\' successfully mapped.')
            except KeyError:
                print(f'>> ERROR: Subtitle for \'{targetFileMetadataSummary}\' not found!')
            targetFileName = targetMetadata[targetFileMetadataSummary]['filename']+'.srt'
            shutil.move(os.path.join(self.tempStorageDir, sourceFileName), \
                        os.path.join(self.targetDir, targetFileName))

        self.cleanTempStorageDir(subtitleFileNames)


if __name__ == '__main__':
    subtitleOp = Subtitles(
                        '/Users/st/Downloads',
                        '/Volumes/Untitled/Game of Thrones/Season 2',
                        '/tmp')
    subtitleOp.mapSubtitlesToFiles()

\$\endgroup\$

1 Answer 1

1
\$\begingroup\$

snake_case

The PEP8 standard is for these:

getTargetFileNames
currentWorkingDir 

to be

get_target_filenames
current_working_dir

and so on for your other functions and variable names.

Fail-safe directory reversion

    os.chdir(currentWorkingDir)

should be in a finally, with the try after the first chdir. That said, you shouldn't be changing the working directory at all. Nothing in getTargetFileNames requires it.

pathlib

Use of pathlib.Path in nearly all of your directory and filesystem operations rather than os will make for a much cleaner program.

Generators

       listOfSubtitleFileNames = list()
       for fileName in listOfFileNames:
           if fileName.endswith('.srt'):
               listOfSubtitleFileNames.append(fileName)
               zipObj.extract(fileName, self.tempStorageDir)
    return listOfSubtitleFileNames

first of all, "list of" can usually be dropped from variable names for legibility; subtitle_filenames is enough to know that it's an iterable. Second, list() can be replaced with []. It would be even simpler as a list comprehension:

subtitle_filenames = [
    fn for fn in filenames if fn.endswith('.srt')
]
for fn in subtitle_filenames:
    zip_obj.extract(fn, self.temp_storage_dir)
return subtitle_filenames

String interpolation

    fileMetadataSummary = metadata['title'] + '.S' \
          + str(metadata['season']) + 'E' \
          + str(metadata['episode']) + '.' \
          + metadata['quality']

First, avoid \ line continuation whenever possible. The way to avoid that here is parentheses. Also, this is well-suited to string interpolation:

file_metadata_summary = (
    metadata['title'] +
    f'S{metadata["season"]}'
    f'E{metadata["episode"]}'
    f'.{metadata["quality"]}'
)

Hard-coded paths

                    '/Users/st/Downloads',
                    '/Volumes/Untitled/Game of Thrones/Season 2',
                    '/tmp'

should not be hard-coded. It should be parametric, perhaps as a command-line parameter.

\$\endgroup\$
2
  • \$\begingroup\$ First of all thanks a lot for your code review. I've made the indicated changes and the code looks much cleaner now. However, I have a question. After making changes according to the GENERATORS section of your answer, my code runs slightly faster. Why using two for loops (1 list comprehension loop and 1 regular loop) is faster than a for loop with an if statement inside it? My intuition says the latter should run faster. It would be a great help if you could explain this to me. \$\endgroup\$ Jun 24, 2020 at 14:00
  • \$\begingroup\$ It's tough to say, and the only concrete answer is going to come from profiling. \$\endgroup\$
    – Reinderien
    Jun 24, 2020 at 14:04

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.