diff options
| -rw-r--r-- | libs/subliminal_patch/providers/utils.py | 55 | ||||
| -rw-r--r-- | tests/subliminal_patch/test_utils.py | 8 |
2 files changed, 53 insertions, 10 deletions
diff --git a/libs/subliminal_patch/providers/utils.py b/libs/subliminal_patch/providers/utils.py index bf3aa856b..ea8e411e6 100644 --- a/libs/subliminal_patch/providers/utils.py +++ b/libs/subliminal_patch/providers/utils.py @@ -4,9 +4,12 @@ import io import logging import os import re +import tempfile +from typing import Iterable, Union import zipfile from guessit import guessit +import pysubs2 import rarfile from subliminal.subtitle import fix_line_ending from subliminal_patch.core import Episode @@ -119,10 +122,10 @@ def is_episode(content): def get_archive_from_bytes(content: bytes): - """Get RarFile/ZipFile object from bytes. Return None is something else - is found.""" - # open the archive + """Get RarFile/ZipFile object from bytes. A ZipFile instance will be returned + if a subtitle-like stream is found. Return None if something else is found.""" archive_stream = io.BytesIO(content) + if rarfile.is_rarfile(archive_stream): logger.debug("Identified rar archive") return rarfile.RarFile(archive_stream) @@ -130,18 +133,50 @@ def get_archive_from_bytes(content: bytes): logger.debug("Identified zip archive") return zipfile.ZipFile(archive_stream) - logger.debug("Unknown compression format") + logger.debug("No compression format found. Trying with subtitle-like files") + + # If the file is a subtitle-like file + with tempfile.NamedTemporaryFile(prefix="spsub", suffix=".srt") as tmp_f: + try: + tmp_f.write(content) + sub = pysubs2.load(tmp_f.name) + except Exception as error: + logger.debug("Couldn't load file: '%s'", error) + else: + if sub is not None: + logger.debug("Identified subtitle file: %s", sub) + zip_obj = zipfile.ZipFile(io.BytesIO(), mode="x") + zip_obj.write(tmp_f.name, os.path.basename(tmp_f.name)) + return zip_obj + + logger.debug("Nothing found") return None -def update_matches(matches, video, release_info: str, **guessit_options): - "Update matches set from release info string. New lines are iterated." +def update_matches( + matches, + video, + release_info: Union[str, Iterable[str]], + split="\n", + **guessit_options +): + """Update matches set from release info string or Iterable. + + Use the split parameter to iterate over the set delimiter; set None to avoid split.""" + 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) + if isinstance(release_info, str): + release_info = release_info.split(split) + + for release in release_info: + for release_split in release.split(split): + logger.debug("Updating matches from release info: %s", release) + matches |= guess_matches( + video, guessit(release_split.strip(), guessit_options) + ) + logger.debug("New matches: %s", matches) return matches diff --git a/tests/subliminal_patch/test_utils.py b/tests/subliminal_patch/test_utils.py index eb8dc2421..c90cf972a 100644 --- a/tests/subliminal_patch/test_utils.py +++ b/tests/subliminal_patch/test_utils.py @@ -122,6 +122,14 @@ def test_update_matches(movies): assert "source" in matches +def test_update_matches_iterable(movies): + matches = set() + utils.update_matches( + matches, movies["dune"], ["Subs for dune 2021 bluray x264", "Dune webrip x264"] + ) + assert "source" in matches + + @pytest.mark.parametrize( "content,expected", [("the.wire.s01e01", True), ("taxi driver 1976", False)] ) |