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