summaryrefslogtreecommitdiffhomepage
path: root/custom_libs/subliminal_patch/subtitle.py
blob: 36afea6d23e3a85e274dd806e59fbcad48abfdbd (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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
# coding=utf-8


from __future__ import absolute_import
import logging
import traceback

import re
import types

import chardet
import pysrt
import pysubs2
from bs4 import UnicodeDammit
from pysubs2 import SSAStyle
from pysubs2.formats.subrip import parse_tags, MAX_REPRESENTABLE_TIME
from pysubs2.time import ms_to_times
from subzero.modification import SubtitleModifications
from subzero.language import Language
from subliminal import Subtitle as Subtitle_
from subliminal.subtitle import Episode, Movie, sanitize_release_group, get_equivalent_release_groups
from subliminal_patch.utils import sanitize
from ftfy import fix_text
from codecs import BOM_UTF8, BOM_UTF16_BE, BOM_UTF16_LE, BOM_UTF32_BE, BOM_UTF32_LE
from six import text_type

BOMS = (
    (BOM_UTF8, "UTF-8"),
    (BOM_UTF32_BE, "UTF-32-BE"),
    (BOM_UTF32_LE, "UTF-32-LE"),
    (BOM_UTF16_BE, "UTF-16-BE"),
    (BOM_UTF16_LE, "UTF-16-LE"),
)

logger = logging.getLogger(__name__)


ftfy_defaults = {
    "uncurl_quotes": False,
    "fix_character_width": False,
}


class Subtitle(Subtitle_):
    storage_path = None
    release_info = None
    matches = {}
    hash_verifiable = False
    hearing_impaired_verifiable = False
    mods = None
    plex_media_fps = None
    skip_wrong_fps = False
    wrong_fps = False
    wrong_series = False
    wrong_season_ep = False
    is_pack = False
    asked_for_release_group = None
    asked_for_episode = None
    uploader = None # string - uploader username

    pack_data = None
    _guessed_encoding = None
    _is_valid = False
    use_original_format = False
    # format = "srt" # default format is srt

    def __init__(self, language, hearing_impaired=False, page_link=None, encoding=None, mods=None, original_format=False):
        # set subtitle language to hi if it's hearing_impaired
        if hearing_impaired:
            language = Language.rebuild(language, hi=True)

        super(Subtitle, self).__init__(language, hearing_impaired=hearing_impaired, page_link=page_link,
                                       encoding=encoding)
        self.mods = mods
        self._is_valid = False
        self.use_original_format = original_format
        self._og_format = None

    @property
    def format(self):
        if self.use_original_format and self._og_format is not None:
            logger.debug("Original format requested [%s]", self._og_format)
            return self._og_format

        logger.debug("Will assume srt format")
        return "srt"

    # Compatibility
    @format.setter
    def format(self, val):
        self._og_format = val

    def __repr__(self):
        r_info = str(self.release_info or "").replace("\n", " | ").strip()
        return f"<{self.__class__.__name__}: {r_info} [{repr(self.language)}]>"

    @property
    def text(self):
        """Content as string

        If :attr:`encoding` is None, the encoding is guessed with :meth:`guess_encoding`

        """
        if not self.content:
            return

        if not isinstance(self.content, text_type):
            return self.content.decode(self.get_encoding(), errors='replace')

        return self.content

    @property
    def numeric_id(self):
        raise NotImplemented

    def get_fps(self):
        """
        :return: frames per second or None if not supported
        :rtype: float
        """
        return None

    def make_picklable(self):
        """
        some subtitle instances might have unpicklable objects stored; clean them up here
        :return: self
        """
        return self

    def get_encoding(self):
        return self.guess_encoding()

    def set_encoding(self, encoding):
        ge = self.get_encoding()
        if encoding == ge:
            return

        unicontent = self.text
        logger.debug("Changing encoding: to %s, from %s", encoding, ge)
        self.content = unicontent.encode(encoding)
        self._guessed_encoding = encoding

    def normalize(self):
        """
        Set encoding to UTF-8 and normalize line endings
        :return:
        """
        self.set_encoding("utf-8")

        # normalize line endings
        self.content = self.content.replace(b"\r\n", b"\n").replace(b'\r', b'\n')

    def _check_bom(self, data):
        return [encoding for bom, encoding in BOMS if data.startswith(bom)]

    def guess_encoding(self):
        """Guess encoding using the language, falling back on chardet.

        :return: the guessed encoding.
        :rtype: str

        """
        if self._guessed_encoding:
            return self._guessed_encoding

        if self.encoding:
            # check provider encoding and use it only if it is valid
            try:
                self.content.decode(self.encoding)
                self._guessed_encoding = self.encoding
                return self._guessed_encoding
            except:
                # provider specified encoding is invalid, fallback to guessing
                pass

        logger.info('Guessing encoding for language %s', self.language)

        encodings = ['utf-8']

        # check UTF BOMs
        bom_encodings = self._check_bom(self.content)
        if bom_encodings:
            encodings = list(set(enc.lower() for enc in bom_encodings + encodings))

        # add language-specific encodings
        # http://scratchpad.wikia.com/wiki/Character_Encoding_Recommendation_for_Languages

        if self.language.alpha3 == 'zho':
            encodings.extend(['cp936', 'gb2312', 'gbk', 'hz', 'iso2022_jp_2', 'cp950', 'big5hkscs', 'big5',
                              'gb18030', 'utf-16'])
        elif self.language.alpha3 == 'jpn':
            encodings.extend(['shift-jis', 'cp932', 'euc_jp', 'iso2022_jp', 'iso2022_jp_1', 'iso2022_jp_2',
                              'iso2022_jp_2004', 'iso2022_jp_3', 'iso2022_jp_ext', ])
        elif self.language.alpha3 == 'tha':
            encodings.extend(['tis-620', 'cp874'])

        # arabian/farsi
        elif self.language.alpha3 in ('ara', 'fas', 'per'):
            encodings.extend(['windows-1256', 'utf-16', 'utf-16le', 'ascii', 'iso-8859-6'])
        elif self.language.alpha3 == 'heb':
            encodings.extend(['windows-1255', 'iso-8859-8'])
        elif self.language.alpha3 == 'tur':
            encodings.extend(['windows-1254', 'iso-8859-9', 'iso-8859-3'])

        # Greek
        elif self.language.alpha3 in ('grc', 'gre', 'ell'):
            encodings.extend(['windows-1253', 'cp1253', 'cp737', 'iso8859-7', 'cp875', 'cp869', 'iso2022_jp_2',
                              'mac_greek'])

        # Polish, Czech, Slovak, Hungarian, Slovene, Bosnian, Croatian, Serbian (Latin script),
        # Romanian and Albanian
        elif self.language.alpha3 in ('pol', 'cze', 'ces', 'slk', 'slo', 'slv', 'hun', 'bos', 'hbs', 'hrv', 'rsb',
                                      'ron', 'rum', 'sqi', 'alb'):

            encodings.extend(['windows-1250', 'iso-8859-2'])

            # Eastern European Group 1
            if self.language.alpha3 == "slv":
                encodings.append('iso-8859-4')

            # Albanian
            elif self.language.alpha3 in ("sqi", "alb"):
                encodings.extend(['windows-1252', 'iso-8859-15', 'iso-8859-1', 'iso-8859-9'])

        # Bulgarian, Serbian and Macedonian, Ukranian and Russian
        elif self.language.alpha3 in ('bul', 'srp', 'mkd', 'mac', 'rus', 'ukr'):
            # Eastern European Group 2
            if self.language.alpha3 in ('bul', 'mkd', 'mac', 'rus', 'ukr'):
                encodings.extend(['windows-1251', 'iso-8859-5'])

            elif self.language.alpha3 == 'srp':
                if self.language.script == "Latn":
                    encodings.extend(['windows-1250', 'iso-8859-2'])
                elif self.language.script == "Cyrl":
                    encodings.extend(['windows-1251', 'iso-8859-5'])
                else:
                    encodings.extend(['windows-1250', 'windows-1251', 'iso-8859-2', 'iso-8859-5'])

        else:
            # Western European (windows-1252) / Northern European
            encodings.extend(['windows-1252', 'iso-8859-15', 'iso-8859-9', 'iso-8859-4', 'iso-8859-1'])

        # try to decode
        logger.debug('Trying encodings %r', encodings)
        for encoding in encodings:
            try:
                self.content.decode(encoding)

            except UnicodeDecodeError:
                pass
            else:
                logger.info('Guessed encoding %s', encoding)
                self._guessed_encoding = encoding
                return encoding

        logger.warning('Could not guess encoding from language')

        # fallback on chardet
        encoding = chardet.detect(self.content)['encoding']
        logger.info('Chardet found encoding %s', encoding)

        if not encoding:
            # fallback on bs4
            logger.info('Falling back to bs4 detection')
            a = UnicodeDammit(self.content)

            logger.info("bs4 detected encoding: %s", a.original_encoding)

            if a.original_encoding:
                self._guessed_encoding = a.original_encoding
                return a.original_encoding
            raise ValueError(u"Couldn't guess the proper encoding for %s", self)

        self._guessed_encoding = encoding
        return encoding

    def is_valid(self):
        """Check if a :attr:`text` is a valid SubRip format. Note that orignal format will pypass the checking

        :return: whether or not the subtitle is valid.
        :rtype: bool

        """
        if self._is_valid:
            return True

        text = self.text
        if not text:
            return False

        # valid srt
        try:
            pysrt.from_string(text, error_handling=pysrt.ERROR_RAISE)
        except Exception:
            logger.error("PySRT-parsing failed, trying pysubs2")
        else:
            self._is_valid = True
            return True

        # something else, try to return srt
        try:
            logger.debug("Trying parsing with PySubs2")
            try:
                # in case of microdvd, try parsing the fps from the subtitle
                subs = pysubs2.SSAFile.from_string(text)
                if subs.format == "microdvd":
                    logger.info("Got FPS from MicroDVD subtitle: %s", subs.fps)
                else:
                    logger.info("Got format: %s", subs.format)
                    self._og_format = subs.format
                    self._is_valid = True
                    # if self.use_original_format:
                    #    self.format = subs.format
                    #    self._is_valid = True
                    #    logger.debug("Using original format")
                    return True

            except pysubs2.UnknownFPSError:
                # if parsing failed, use frame rate from provider
                sub_fps = self.get_fps()
                if not isinstance(sub_fps, float) or sub_fps < 10.0:
                    # or use our media file's fps as a fallback
                    sub_fps = self.plex_media_fps
                    logger.info("No FPS info in subtitle. Using our own media FPS for the MicroDVD subtitle: %s",
                                self.plex_media_fps)
                subs = pysubs2.SSAFile.from_string(text, fps=sub_fps)

            unicontent = self.pysubs2_to_unicode(subs)
            self.content = unicontent.encode(self.get_encoding())
        except:
            logger.exception("Couldn't convert subtitle %s to .srt format: %s", self, traceback.format_exc())
            return False

        self._is_valid = True
        return True

    @classmethod
    def pysubs2_to_unicode(cls, sub, format="srt"):
        """
        this is a modified version of pysubs2.SubripFormat.to_file with special handling for drawing tags in ASS
        :param sub:
        :param format:
        :return:
        """
        def ms_to_timestamp(ms, mssep=","):
            """Convert ms to 'HH:MM:SS,mmm'"""
            # XXX throw on overflow/underflow?
            if ms < 0: ms = 0
            if ms > MAX_REPRESENTABLE_TIME: ms = MAX_REPRESENTABLE_TIME
            h, m, s, ms = ms_to_times(ms)
            return "%02d:%02d:%02d%s%03d" % (h, m, s, mssep, ms)

        def prepare_text(text, style):
            body = []
            for fragment, sty in parse_tags(text, style, sub.styles):
                fragment = fragment.replace(r"\h", u" ")
                fragment = fragment.replace(r"\n", u"\n")
                fragment = fragment.replace(r"\N", u"\n")
                if sty.drawing:
                    raise pysubs2.ContentNotUsable

                if format == "srt":
                    if sty.italic:
                        fragment = u"<i>%s</i>" % fragment
                    if sty.underline:
                        fragment = u"<u>%s</u>" % fragment
                    if sty.strikeout:
                        fragment = u"<s>%s</s>" % fragment
                elif format == "vtt":
                    if sty.bold:
                        fragment = u"<b>%s</b>" % fragment
                    if sty.italic:
                        fragment = u"<i>%s</i>" % fragment
                    if sty.underline:
                        fragment = u"<u>%s</u>" % fragment

                body.append(fragment)

            return re.sub(u"\n+", u"\n", u"".join(body).strip())

        visible_lines = (line for line in sub if not line.is_comment)

        out = []
        mssep = ","

        if format == "vtt":
            out.append("WEBVTT\n\n")
            mssep = "."

        for i, line in enumerate(visible_lines, 1):
            start = ms_to_timestamp(line.start, mssep=mssep)
            end = ms_to_timestamp(line.end, mssep=mssep)
            try:
                text = prepare_text(line.text, sub.styles.get(line.style, SSAStyle.DEFAULT_STYLE))
            except pysubs2.ContentNotUsable:
                continue

            out.append(u"%d\n" % i)
            out.append(u"%s --> %s\n" % (start, end))
            out.append(u"%s%s" % (text, "\n\n"))

        return u"".join(out)

    def get_modified_content(self, format="srt", debug=False):
        """
        :return: string
        """
        if not self.mods:
            return fix_text(self.content.decode(encoding=self.get_encoding()), **ftfy_defaults).encode(
                encoding=self.get_encoding())

        submods = SubtitleModifications(debug=debug)
        if submods.load(content=self.text, language=self.language):
            logger.info("Applying mods: %s", self.mods)
            submods.modify(*self.mods)
            self.mods = submods.mods_used

            content = fix_text(self.pysubs2_to_unicode(submods.f, format=format), **ftfy_defaults)\
                .encode(encoding=self.get_encoding())
            submods.f = None
            del submods
            return content
        return None


class ModifiedSubtitle(Subtitle):
    id = None


MERGED_FORMATS = {
    "TV": ("HDTV", "SDTV", "AHDTV", "Ultra HDTV"),
    "Air": ("SATRip", "DVB", "PPV", "Digital TV"),
    "Disk-HD": ("HD-DVD", "Blu-ray", "Ultra HD Blu-ray"),
    "Disk-SD": ("DVD", "VHS"),
    "Web": ("Web",),
}

MERGED_FORMATS_REV = dict((v.lower(), k.lower()) for k in MERGED_FORMATS for v in MERGED_FORMATS[k])

def _has_match(video, guess, key) -> bool:
    value = getattr(video, key)
    guess_value = guess.get(key)

    # To avoid extra debug calls
    if guess_value is None or value is None:
        return False

    if isinstance(guess_value, list):
        matched = any(value == item for item in guess_value)
    else:
        matched = value == guess_value

    logger.debug("%s matched? %s (%s -> %s)", key, matched, value, guess_value)

    return matched



def guess_matches(video, guess, partial=False):
    """Get matches between a `video` and a `guess`.

    If a guess is `partial`, the absence information won't be counted as a match.

    Patch: add multiple release group and formats handling

    :param video: the video.
    :type video: :class:`~subliminal.video.Video`
    :param guess: the guess.
    :type guess: dict
    :param bool partial: whether or not the guess is partial.
    :return: matches between the `video` and the `guess`.
    :rtype: set

    """
    matches = set()
    if isinstance(video, Episode):
        # series
        if video.series and 'title' in guess:
            titles = guess["title"]
            if not isinstance(titles, list):
                titles = [titles]

            for title in titles:
                if sanitize(title) in (sanitize(name) for name in [video.series] + video.alternative_series):
                    matches.add('series')

        # title
        if video.title and 'episode_title' in guess and sanitize(guess['episode_title']) == sanitize(video.title):
            matches.add('title')

        # season
        if video.season and 'season' in guess and guess['season'] == video.season:
            matches.add('season')

        # episode
        # Currently we only have single-ep support (guessit returns a multi-ep as a list with int values)
        # Most providers only support single-ep, so make sure it contains only 1 episode
        # In case of multi-ep, take the lowest episode (subtitles will normally be available on lowest episode number)
        if video.episode and 'episode' in guess:
            episode = episode_guess = guess['episode']
            if isinstance(episode_guess, list):
                try:
                    episode = min([int(x) for x in episode_guess])
                except (TypeError, ValueError):
                    pass
            if episode == video.episode:
                matches.add('episode')

        # year
        if video.year and 'year' in guess and guess['year'] == video.year:
            matches.add('year')

        # count "no year" as an information
        if not partial and video.original_series and 'year' not in guess:
            matches.add('year')

    elif isinstance(video, Movie):
        # year
        if video.year and 'year' in guess and guess['year'] == video.year:
            matches.add('year')
        # title
        if video.title and 'title' in guess and sanitize(guess['title']) in (
                    sanitize(name) for name in [video.title] + video.alternative_titles):
            matches.add('title')

    # release_group
    if 'release_group' in guess:
        release_groups = guess["release_group"]
        if not isinstance(release_groups, list):
            release_groups = [release_groups]

        if video.release_group:
            for release_group in release_groups:
                if (sanitize_release_group(release_group) in
                        get_equivalent_release_groups(sanitize_release_group(video.release_group))):
                    matches.add('release_group')
                    break
    # source
    if 'source' in guess:
        formats = guess["source"]
        if not isinstance(formats, list):
            formats = [formats]

        if video.source:
            video_format = video.source.lower()
            _video_gen_format = MERGED_FORMATS_REV.get(video_format)
            matched = False
            for frmt in formats:
                _guess_gen_frmt = MERGED_FORMATS_REV.get(frmt.lower())
                # We don't want to match a singleton
                if _guess_gen_frmt is None: # If the source is not in MERGED_FORMATS
                    _guess_gen_frmt = guess["source"]

                if _guess_gen_frmt == _video_gen_format:
                    matched = True
                    matches.add('source')
                    break

            logger.debug("Source match found? %s: %s -> %s", matched, video.source, formats)

        if "release_group" in matches and "source" not in matches:
            logger.info("Release group matched but source didn't. Removing release group match.")
            matches.remove("release_group")

    guess.update({"resolution": guess.get("screen_size")})

    # Solve match keys for potential lists
    for key in ("video_codec", "audio_codec", "edition", "streaming_service", "resolution"):
        if _has_match(video, guess, key):
            matches.add(key)

    for key in ("streaming_service", "edition", "other"):
        if _check_optional(video, guess, key):
            matches.add(key)

    return matches


def _check_optional(video, guess, key="edition"):
    guess_optional = guess.get(key)
    video_optional = getattr(video, key, None)

    if video_optional and guess_optional:
        return _has_match(video, guess, key)

    if not video_optional and not guess_optional:
        logger.debug("Both video and guess don't have %s. Returning True", key)
        return True

    logger.debug("One item doesn't have %s (%s -> %s). Returning False", key, guess_optional, video_optional)
    return False