1import copy
2import functools
3import re
4from typing import Dict, Iterator, List, Optional, Set, Tuple, Union
5
6from loguru import logger
7
8from flexget.utils.serialization import Serializer
9
10logger = logger.bind(name='utils.qualities')
11
12
13@functools.total_ordering
14class QualityComponent:
15    """"""
16
17    def __init__(
18        self,
19        type: str,
20        value: int,
21        name: str,
22        regexp: Optional[str] = None,
23        modifier: Optional[int] = None,
24        defaults: Optional[List['QualityComponent']] = None,
25    ) -> None:
26        """
27        :param type: Type of quality component. (resolution, source, codec, color_range or audio)
28        :param value: Value used to sort this component with others of like type.
29        :param name: Canonical name for this quality component.
30        :param regexp: Regexps used to match this component.
31        :param modifier: An integer that affects sorting above all other components.
32        :param defaults: An iterable defining defaults for other quality components if this component matches.
33        """
34
35        if type not in ['resolution', 'source', 'codec', 'color_range', 'audio']:
36            raise ValueError('%s is not a valid quality component type.' % type)
37        self.type = type
38        self.value = value
39        self.name = name
40        self.modifier = modifier
41        self.defaults = defaults or []
42
43        # compile regexp
44        if regexp is None:
45            regexp = re.escape(name)
46        self.regexp = re.compile(r'(?<![^\W_])(' + regexp + r')(?![^\W_])', re.IGNORECASE)
47
48    def matches(self, text: str) -> Tuple[bool, str]:
49        """Test if quality matches to text.
50
51        :param string text: data te be tested against
52        :returns: tuple (matches, remaining text without quality data)
53        """
54
55        match = self.regexp.search(text)
56        if not match:
57            return False, ""
58        else:
59            # remove matching part from the text
60            text = text[: match.start()] + text[match.end() :]
61        return True, text
62
63    def __hash__(self) -> int:
64        return hash(self.type + str(self.value))
65
66    def __bool__(self) -> bool:
67        return bool(self.value)
68
69    def __eq__(self, other) -> bool:
70        if isinstance(other, str):
71            other = _registry.get(other)
72        if not isinstance(other, QualityComponent):
73            raise TypeError('Cannot compare %r and %r' % (self, other))
74        if other.type == self.type:
75            return self.value == other.value
76        else:
77            raise TypeError('Cannot compare %s and %s' % (self.type, other.type))
78
79    def __lt__(self, other) -> bool:
80        if isinstance(other, str):
81            other = _registry.get(other)
82        if not isinstance(other, QualityComponent):
83            raise TypeError('Cannot compare %r and %r' % (self, other))
84        if other.type == self.type:
85            return self.value < other.value
86        else:
87            raise TypeError('Cannot compare %s and %s' % (self.type, other.type))
88
89    def __add__(self, other):
90        if not isinstance(other, int):
91            raise TypeError()
92        l = globals().get('_' + self.type + 's')
93        index = l.index(self) + other
94        if index >= len(l):
95            index = -1
96        return l[index]
97
98    def __sub__(self, other):
99        if not isinstance(other, int):
100            raise TypeError()
101        l = globals().get('_' + self.type + 's')
102        index = l.index(self) - other
103        if index < 0:
104            index = 0
105        return l[index]
106
107    def __repr__(self) -> str:
108        return f'<{self.type.title()}(name={self.name},value={self.value})>'
109
110    def __str__(self) -> str:
111        return self.name
112
113    def __deepcopy__(self, memo=None):
114        # No mutable attributes, return a regular copy
115        return copy.copy(self)
116
117
118_resolutions = [
119    QualityComponent('resolution', 10, '360p'),
120    QualityComponent('resolution', 20, '368p', '368p?'),
121    QualityComponent('resolution', 30, '480p', '480p?'),
122    QualityComponent('resolution', 35, '540p', '540p?'),
123    QualityComponent('resolution', 40, '576p', '576p?'),
124    QualityComponent('resolution', 45, 'hr'),
125    QualityComponent('resolution', 50, '720i'),
126    QualityComponent('resolution', 60, '720p', '(1280x)?720(p|hd)?x?([56]0)?'),
127    QualityComponent('resolution', 70, '1080i'),
128    QualityComponent('resolution', 80, '1080p', '(1920x)?1080p?x?([56]0)?'),
129    QualityComponent('resolution', 90, '2160p', '((3840x)?2160p?x?([56]0)?)|4k'),
130]
131_sources = [
132    QualityComponent('source', 10, 'workprint', modifier=-8),
133    QualityComponent('source', 20, 'cam', '(?:hd)?cam', modifier=-7),
134    QualityComponent('source', 30, 'ts', '(?:hd)?ts|telesync', modifier=-6),
135    QualityComponent('source', 40, 'tc', 'tc|telecine', modifier=-5),
136    QualityComponent('source', 50, 'r5', 'r[2-8c]', modifier=-4),
137    QualityComponent('source', 60, 'hdrip', r'hd[\W_]?rip', modifier=-3),
138    QualityComponent('source', 70, 'ppvrip', r'ppv[\W_]?rip', modifier=-2),
139    QualityComponent('source', 80, 'preair', modifier=-1),
140    QualityComponent('source', 90, 'tvrip', r'tv[\W_]?rip'),
141    QualityComponent('source', 100, 'dsr', r'dsr|ds[\W_]?rip'),
142    QualityComponent('source', 110, 'sdtv', r'(?:[sp]dtv|dvb)(?:[\W_]?rip)?'),
143    QualityComponent('source', 120, 'dvdscr', r'(?:(?:dvd|web)[\W_]?)?scr(?:eener)?', modifier=0),
144    QualityComponent('source', 130, 'bdscr', 'bdscr(?:eener)?'),
145    QualityComponent('source', 140, 'webrip', r'web[\W_]?rip'),
146    QualityComponent('source', 150, 'hdtv', r'a?hdtv(?:[\W_]?rip)?'),
147    QualityComponent('source', 160, 'webdl', r'web(?:[\W_]?(dl|hd))?'),
148    QualityComponent('source', 170, 'dvdrip', r'dvd(?:[\W_]?rip)?'),
149    QualityComponent('source', 175, 'remux'),
150    QualityComponent('source', 180, 'bluray', r'(?:b[dr][\W_]?rip|blu[\W_]?ray(?:[\W_]?rip)?)'),
151]
152_codecs = [
153    QualityComponent('codec', 10, 'divx'),
154    QualityComponent('codec', 20, 'xvid'),
155    QualityComponent('codec', 30, 'h264', '[hx].?264'),
156    QualityComponent('codec', 35, 'vp9'),
157    QualityComponent('codec', 40, 'h265', '[hx].?265|hevc'),
158]
159
160_color_ranges = [
161    QualityComponent('color_range', 10, '8bit', r'8[^\w]?bits?|hi8p?'),
162    QualityComponent('color_range', 20, '10bit', r'10[^\w]?bits?|hi10p?'),
163    QualityComponent('color_range', 40, 'hdrplus', r'hdr[^\w]?(\+|p|plus)'),
164    QualityComponent('color_range', 30, 'hdr', r'hdr([^\w]?10)?'),
165    QualityComponent('color_range', 50, 'dolbyvision', r'(dolby[^\w]?vision|dv|dovi)'),
166]
167
168channels = r'(?:(?:[^\w+]?[1-7][\W_]?(?:0|1|ch)))'
169_audios = [
170    QualityComponent('audio', 10, 'mp3'),
171    # TODO: No idea what order these should go in or if we need different regexps
172    QualityComponent('audio', 20, 'aac', 'aac%s?' % channels),
173    QualityComponent('audio', 30, 'dd5.1', 'dd%s' % channels),
174    QualityComponent('audio', 40, 'ac3', 'ac3%s?' % channels),
175    QualityComponent('audio', 45, 'dd+5.1', 'dd[p+]%s' % channels),
176    QualityComponent('audio', 50, 'flac', 'flac%s?' % channels),
177    # The DTSs are a bit backwards, but the more specific one needs to be parsed first
178    QualityComponent('audio', 70, 'dtshd', r'dts[\W_]?hd(?:[\W_]?ma)?%s?' % channels),
179    QualityComponent('audio', 60, 'dts'),
180    QualityComponent('audio', 80, 'truehd', 'truehd%s?' % channels),
181]
182
183_UNKNOWNS = {
184    'resolution': QualityComponent('resolution', 0, 'unknown'),
185    'source': QualityComponent('source', 0, 'unknown'),
186    'codec': QualityComponent('codec', 0, 'unknown'),
187    'color_range': QualityComponent('color_range', 0, 'unknown'),
188    'audio': QualityComponent('audio', 0, 'unknown'),
189}
190
191# For wiki generating help
192'''for type in (_resolutions, _sources, _codecs, _color_ranges, _audios):
193    print '{{{#!td style="vertical-align: top"'
194    for item in reversed(type):
195        print '- ' + item.name
196    print '}}}'
197'''
198
199_registry: Dict[Union[str, QualityComponent], QualityComponent] = {}
200for items in (_resolutions, _sources, _codecs, _color_ranges, _audios):
201    for item in items:
202        _registry[item.name] = item
203
204
205def all_components() -> Iterator[QualityComponent]:
206    return iter(_registry.values())
207
208
209@functools.total_ordering
210class Quality(Serializer):
211    """Parses and stores the quality of an entry in the four component categories."""
212
213    def __init__(self, text: str = '') -> None:
214        """
215        :param text: A string to parse quality from
216        """
217        self.text = text
218        self.clean_text = text
219        if text:
220            self.parse(text)
221        else:
222            self.resolution = _UNKNOWNS['resolution']
223            self.source = _UNKNOWNS['source']
224            self.codec = _UNKNOWNS['codec']
225            self.color_range = _UNKNOWNS['color_range']
226            self.audio = _UNKNOWNS['audio']
227
228    def parse(self, text: str) -> None:
229        """Parses a string to determine the quality in the four component categories.
230
231        :param text: The string to parse
232        """
233        self.text = text
234        self.clean_text = text
235        self.resolution = self._find_best(_resolutions, _UNKNOWNS['resolution'], False)
236        self.source = self._find_best(_sources, _UNKNOWNS['source'])
237        self.codec = self._find_best(_codecs, _UNKNOWNS['codec'])
238        self.color_range = self._find_best(_color_ranges, _UNKNOWNS['color_range'])
239        self.audio = self._find_best(_audios, _UNKNOWNS['audio'])
240        # If any of the matched components have defaults, set them now.
241        for component in self.components:
242            for default in component.defaults:
243                default = _registry[default]
244                if not getattr(self, default.type):
245                    setattr(self, default.type, default)
246
247    def _find_best(
248        self,
249        qlist: List[QualityComponent],
250        default: QualityComponent,
251        strip_all: bool = True,
252    ) -> QualityComponent:
253        """Finds the highest matching quality component from `qlist`"""
254        result = None
255        search_in = self.clean_text
256        for item in qlist:
257            match = item.matches(search_in)
258            if match[0]:
259                result = item
260                self.clean_text = match[1]
261                if strip_all:
262                    # In some cases we want to strip all found quality components,
263                    # even though we're going to return only the last of them.
264                    search_in = self.clean_text
265                if item.modifier is not None:
266                    # If this item has a modifier, do not proceed to check higher qualities in the list
267                    break
268        return result or default
269
270    @property
271    def name(self) -> str:
272        name = ' '.join(
273            str(p)
274            for p in (self.resolution, self.source, self.codec, self.color_range, self.audio)
275            if p.value != 0
276        )
277        return name or 'unknown'
278
279    @property
280    def components(self) -> List[QualityComponent]:
281        return [self.resolution, self.source, self.codec, self.color_range, self.audio]
282
283    @classmethod
284    def serialize(cls, quality: 'Quality') -> str:
285        return str(quality)
286
287    @classmethod
288    def deserialize(cls, data: str, version: int) -> 'Quality':
289        return cls(data)
290
291    @property
292    def _comparator(self) -> List:
293        modifier = sum(c.modifier for c in self.components if c.modifier)
294        return [modifier] + self.components
295
296    def __contains__(self, other):
297        if isinstance(other, str):
298            other = Quality(other)
299        if not other or not self:
300            return False
301        for cat in ('resolution', 'source', 'audio', 'color_range', 'codec'):
302            othercat = getattr(other, cat)
303            if othercat and othercat != getattr(self, cat):
304                return False
305        return True
306
307    def __bool__(self) -> bool:
308        return any(self._comparator)
309
310    def __eq__(self, other) -> bool:
311        if isinstance(other, str):
312            other = Quality(other)
313        if not isinstance(other, Quality):
314            if other is None:
315                return False
316            raise TypeError('Cannot compare %r and %r' % (self, other))
317        return self._comparator == other._comparator
318
319    def __lt__(self, other) -> bool:
320        if isinstance(other, str):
321            other = Quality(other)
322        if not isinstance(other, Quality):
323            raise TypeError('Cannot compare %r and %r' % (self, other))
324        return self._comparator < other._comparator
325
326    def __repr__(self) -> str:
327        return '<Quality(resolution=%s,source=%s,codec=%s,color_range=%s,audio=%s)>' % (
328            self.resolution,
329            self.source,
330            self.codec,
331            self.color_range,
332            self.audio,
333        )
334
335    def __str__(self) -> str:
336        return self.name
337
338    def __hash__(self) -> int:
339        # Make these usable as dict keys
340        return hash(self.name)
341
342
343def get(quality_name: str) -> Quality:
344    """Returns a quality object based on canonical quality name."""
345
346    found_components = {}
347    for part in quality_name.lower().split():
348        component = _registry.get(part)
349        if not component:
350            raise ValueError('`%s` is not a valid quality string' % part)
351        if component.type in found_components:
352            raise ValueError('`%s` cannot be defined twice in a quality' % component.type)
353        found_components[component.type] = component
354    if not found_components:
355        raise ValueError('No quality specified')
356    result = Quality()
357    for type, component in found_components.items():
358        setattr(result, type, component)
359    return result
360
361
362class RequirementComponent:
363    """Represents requirements for a given component type. Can evaluate whether a given QualityComponent
364    meets those requirements."""
365
366    def __init__(self, type: str) -> None:
367        self.type = type
368        self.min: Optional[QualityComponent] = None
369        self.max: Optional[QualityComponent] = None
370        self.acceptable: Set[QualityComponent] = set()
371        self.none_of: Set[QualityComponent] = set()
372
373    def reset(self) -> None:
374        self.min = None
375        self.max = None
376        self.acceptable = set()
377        self.none_of = set()
378
379    def allows(self, comp: QualityComponent, loose: bool = False) -> bool:
380        if comp.type != self.type:
381            raise TypeError('Cannot compare %r against %s' % (comp, self.type))
382        if comp in self.none_of:
383            return False
384        if loose:
385            return True
386        if comp in self.acceptable:
387            return True
388        if self.min or self.max:
389            if self.min and comp < self.min:
390                return False
391            if self.max and comp > self.max:
392                return False
393            return True
394        if not self.acceptable:
395            return True
396        return False
397
398    def add_requirement(self, text: str) -> None:
399        if '-' in text:
400            min_str, max_str = text.split('-')
401            min_quality, max_quality = _registry[min_str], _registry[max_str]
402            if min_quality.type != max_quality.type != self.type:
403                raise ValueError('Component type mismatch: %s' % text)
404            self.min, self.max = min_quality, max_quality
405        elif '|' in text:
406            req_quals = text.split('|')
407            quals = {_registry[qual] for qual in req_quals}
408            if any(qual.type != self.type for qual in quals):
409                raise ValueError('Component type mismatch: %s' % text)
410            self.acceptable |= quals
411        else:
412            qual = _registry[text.strip('!<>=+')]
413            if qual.type != self.type:
414                raise ValueError('Component type mismatch!')
415            if text in _registry:
416                self.acceptable.add(qual)
417            else:
418                if text[0] == '<':
419                    if text[1] != '=':
420                        qual -= 1
421                    self.max = qual
422                elif text[0] == '>' or text.endswith('+'):
423                    if text[1] != '=' and not text.endswith('+'):
424                        qual += 1
425                    self.min = qual
426                elif text[0] == '!':
427                    self.none_of.add(qual)
428
429    def __eq__(self, other) -> bool:
430        if not isinstance(other, RequirementComponent):
431            return False
432        return (self.max, self.max, self.acceptable, self.none_of) == (
433            other.max,
434            other.max,
435            other.acceptable,
436            other.none_of,
437        )
438
439    def __hash__(self) -> int:
440        return hash(
441            tuple(
442                [self.min, self.max, tuple(sorted(self.acceptable)), tuple(sorted(self.none_of))]
443            )
444        )
445
446
447class Requirements:
448    """Represents requirements for allowable qualities. Can determine whether a given Quality passes requirements."""
449
450    def __init__(self, req: str = '') -> None:
451        self.text = ''
452        self.resolution = RequirementComponent('resolution')
453        self.source = RequirementComponent('source')
454        self.codec = RequirementComponent('codec')
455        self.color_range = RequirementComponent('color_range')
456        self.audio = RequirementComponent('audio')
457        if req:
458            self.parse_requirements(req)
459
460    @property
461    def components(self) -> List[RequirementComponent]:
462        return [self.resolution, self.source, self.codec, self.color_range, self.audio]
463
464    def parse_requirements(self, text: str) -> None:
465        """
466        Parses a requirements string.
467
468        :param text: The string containing quality requirements.
469        """
470        text = text.lower()
471        if self.text:
472            self.text += ' '
473        self.text += text
474        if self.text == 'any':
475            for component in self.components:
476                component.reset()
477                return
478
479        text = text.replace(',', ' ')
480        parts = text.split()
481        try:
482            for part in parts:
483                if '-' in part:
484                    found = _registry[part.split('-')[0]]
485                elif '|' in part:
486                    found = _registry[part.split('|')[0]]
487                else:
488                    found = _registry[part.strip('!<>=+')]
489                for component in self.components:
490                    if found.type == component.type:
491                        component.add_requirement(part)
492        except KeyError as e:
493            raise ValueError('%s is not a valid quality component.' % e.args[0])
494
495    def allows(self, qual: Union[Quality, str], loose: bool = False) -> bool:
496        """Determine whether this set of requirements allows a given quality.
497
498        :param Quality qual: The quality to evaluate.
499        :param bool loose: If True, only ! (not) requirements will be enforced.
500        :rtype: bool
501        :returns: True if given quality passes all component requirements.
502        """
503        if isinstance(qual, str):
504            qual = Quality(qual)
505        for r_component, q_component in zip(self.components, qual.components):
506            if not r_component.allows(q_component, loose=loose):
507                return False
508        return True
509
510    def __eq__(self, other) -> bool:
511        if isinstance(other, str):
512            other = Requirements(other)
513        return self.components == other.components
514
515    def __hash__(self) -> int:
516        return hash(tuple(self.components))
517
518    def __str__(self) -> str:
519        return self.text or 'any'
520
521    def __repr__(self) -> str:
522        return f'<Requirements({self})>'
523