summaryrefslogtreecommitdiffhomepage
path: root/libs/subliminal_patch/providers/utils.py
blob: 0158abfee831c551782a0b1eb74b71c0f6ce5bd3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import io
import logging
import os
import zipfile

import rarfile

from guessit import guessit

from subliminal.subtitle import fix_line_ending

from subliminal_patch.core import Episode
from subliminal_patch.subtitle import guess_matches

from ._agent_list import FIRST_THOUSAND_OR_SO_USER_AGENTS


logger = logging.getLogger(__name__)


def _get_matching_sub(sub_names, forced=False, episode=None):
    matching_sub = None

    guess_options = {"single_value": True}
    if episode is not None:
        guess_options["type"] = "episode"  # type: ignore

    for sub_name in sub_names:
        if not forced and os.path.splitext(sub_name.lower())[0].endswith("forced"):
            logger.debug("Ignoring forced subtitle: %s", sub_name)
            continue

        # If it's a movie then get the first subtitle
        if episode is None:
            logger.debug("Movie subtitle found: %s", sub_name)
            matching_sub = sub_name
            break

        guess = guessit(sub_name, options=guess_options)

        if guess.get("episode") is None:
            logger.debug("No episode info found in file: %s", sub_name)
            continue

        if episode == guess["episode"]:
            logger.debug("Episode matched: %s", sub_name)
            matching_sub = sub_name
            break

        logger.debug("Ignoring incorrect episode: %s", sub_name)

    return matching_sub


def get_subtitle_from_archive(
    archive, forced=False, episode=None, get_first_subtitle=False
):
    "Get subtitle from Rarfile/Zipfile object. Return None if nothing is found."
    subs_in_archive = [
        name
        for name in archive.namelist()
        if name.endswith((".srt", ".sub", ".ssa", ".ass"))
    ]

    if not subs_in_archive:
        logger.info("No subtitles found in archive")
        return None

    logger.debug("Subtitles in archive: %s", subs_in_archive)

    if len(subs_in_archive) == 1 or get_first_subtitle:
        logger.debug("Getting first subtitle in archive: %s", subs_in_archive)
        return fix_line_ending(archive.read(subs_in_archive[0]))

    matching_sub = _get_matching_sub(subs_in_archive, forced, episode)

    if matching_sub is not None:
        logger.info("Using %s from archive", matching_sub)
        return fix_line_ending(archive.read(matching_sub))

    logger.debug("No subtitle found in archive")
    return None


def is_episode(content):
    return "episode" in guessit(content, {"type": "episode"})


def get_archive_from_bytes(content: bytes):
    """Get RarFile/ZipFile object from bytes. Return None is something else
    is found."""
    # open the archive
    archive_stream = io.BytesIO(content)
    if rarfile.is_rarfile(archive_stream):
        logger.debug("Identified rar archive")
        return rarfile.RarFile(archive_stream)
    elif zipfile.is_zipfile(archive_stream):
        logger.debug("Identified zip archive")
        return zipfile.ZipFile(archive_stream)

    logger.debug("Unknown compression format")
    return None


def update_matches(matches, video, release_info: str, **guessit_options):
    "Update matches set from release info string. New lines are iterated."
    guessit_options["type"] = "episode" if isinstance(video, Episode) else "movie"
    logger.debug("Guessit options to update matches: %s", guessit_options)

    for release in release_info.split("\n"):
        logger.debug("Updating matches from release info: %s", release)
        matches |= guess_matches(video, guessit(release.strip(), guessit_options))
        logger.debug("New matches: %s", matches)

    return matches