From cb3892bd113050d2d88504906fde070cd11567d1 Mon Sep 17 00:00:00 2001
From: Rato <rato.aq2@gmail.com>
Date: Wed, 19 Jul 2017 23:00:35 +0200
Subject: [PATCH] Fix for #276: Added aspect_ratio and validation for standard
 resolutions

---
 guessit/rules/properties/screen_size.py | 102 +++++++------
 guessit/test/rules/screen_size.yml      | 182 +++++++++++++++++++++++-
 2 files changed, 236 insertions(+), 48 deletions(-)

diff --git a/guessit/rules/properties/screen_size.py b/guessit/rules/properties/screen_size.py
index f690c8b..d824ec3 100644
--- a/guessit/rules/properties/screen_size.py
+++ b/guessit/rules/properties/screen_size.py
@@ -3,12 +3,18 @@
 """
 screen_size property
 """
+from rebulk.match import Match
 from rebulk.remodule import re
 
-from rebulk import Rebulk, Rule, RemoveMatch
+from rebulk import Rebulk, Rule, RemoveMatch, AppendMatch
+
+from guessit.reutils import build_or_pattern
 from ..common.validators import seps_surround
 from ..common import dash, seps
 
+interlaced = frozenset({'360', '480', '576', '900', '1080'})
+progressive = frozenset(interlaced | {'368', '720', '2160', '4320'})
+
 
 def screen_size():
     """
@@ -16,56 +22,68 @@ def screen_size():
     :return: Created Rebulk object
     :rtype: Rebulk
     """
-    def conflict_solver(match, other):
-        """
-        Conflict solver for most screen_size.
-        """
-        if other.name == 'screen_size':
-            if 'resolution' in other.tags:
-                # The chtouile to solve conflict in "720 x 432" string matching both 720p pattern
-                int_value = _digits_re.findall(match.raw)[-1]
-                if other.value.startswith(int_value):
-                    return match
-            return other
-        return '__default__'
-
     rebulk = Rebulk().string_defaults(ignore_case=True).regex_defaults(flags=re.IGNORECASE)
-    rebulk.defaults(name="screen_size", validator=seps_surround, conflict_solver=conflict_solver)
 
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?360(?:i)', value='360i')
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?360(?:p?x?)', value='360p')
-    rebulk.regex(r"(?:\d{3,}(?:x|\*))?368(?:p?x?)", value="368p")
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?480(?:i)', value='480i')
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?480(?:p?x?)', value='480p')
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?576(?:i)', value='576i')
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?576(?:p?x?)', value='576p')
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?720(?:p?(?:50|60)?x?)', value='720p')
-    rebulk.regex(r"(?:\d{3,}(?:x|\*))?720(?:p(?:50|60)?x?)", value="720p")
-    rebulk.regex(r"(?:\d{3,}(?:x|\*))?720p?hd", value="720p")
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?900(?:i)', value='900i')
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?900(?:p?x?)', value='900p')
-    rebulk.regex(r"(?:\d{3,}(?:x|\*))?1080i", value="1080i")
-    rebulk.regex(r"(?:\d{3,}(?:x|\*))?1080p?x?", value="1080p")
-    rebulk.regex(r"(?:\d{3,}(?:x|\*))?1080(?:p(?:50|60)?x?)", value="1080p")
-    rebulk.regex(r"(?:\d{3,}(?:x|\*))?1080p?hd", value="1080p")
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?2160(?:p?x?)', value='2160p')
+    rebulk.defaults(name='screen_size', validator=seps_surround, abbreviations=[dash], private_children=True)
+
+    res_pattern = r'(?:(?P<width>\d{3,4})(?:x|\*))?'
+    rebulk.regex(res_pattern + build_or_pattern(interlaced, name='height') + r'(?P<scan_type>i)(?:24|30|50|60|120)?')
+    rebulk.regex(res_pattern + build_or_pattern(progressive, name='height') + r'(?P<scan_type>p)(?:24|30|50|60|120)?')
+    rebulk.regex(res_pattern + build_or_pattern(progressive, name='height') + r'(?P<scan_type>p)?(?:hd)')
+    rebulk.regex(res_pattern + build_or_pattern(progressive, name='height') + r'(?P<scan_type>p)?x?')
     rebulk.string('4k', value='2160p')
-    rebulk.regex(r'(?:\d{3,}(?:x|\*))?4320(?:p?x?)', value='4320p')
-
-    _digits_re = re.compile(r'\d+')
-
-    rebulk.defaults(name="screen_size", validator=seps_surround)
-    rebulk.regex(r'\d{3,}-?(?:x|\*)-?\d{3,}',
-                 formatter=lambda value: 'x'.join(_digits_re.findall(value)),
-                 abbreviations=[dash],
-                 tags=['resolution'],
+    rebulk.regex(r'(?P<width>\d{3,4})-?(?:x|\*)-?(?P<height>\d{3,4})',
                  conflict_solver=lambda match, other: '__default__' if other.name == 'screen_size' else other)
 
-    rebulk.rules(ScreenSizeOnlyOne, RemoveScreenSizeConflicts)
+    rebulk.rules(PostProcessScreenSize(progressive), ScreenSizeOnlyOne, RemoveScreenSizeConflicts)
 
     return rebulk
 
 
+class PostProcessScreenSize(Rule):
+    """
+    Process the screen size calculating the aspect ratio if available.
+
+    Convert to a standard notation (720p, 1080p, etc) when it's a standard resolution and
+    aspect ratio is valid or not available.
+
+    It also creates an aspect_ratio match when available.
+    """
+    consequence = AppendMatch
+
+    def __init__(self, standard_heights, min_ar=1.333, max_ar=1.898):
+        super(PostProcessScreenSize, self).__init__()
+        self.standard_heights = standard_heights
+        self.min_ar = min_ar
+        self.max_ar = max_ar
+
+    def when(self, matches, context):
+        to_append = []
+        for match in matches.named('screen_size'):
+            values = match.children.to_dict()
+            if 'height' not in values:
+                continue
+
+            scan_type = (values.get('scan_type') or 'p').lower()
+            height = values['height']
+            if 'width' not in values:
+                match.value = '{0}{1}'.format(height, scan_type)
+                continue
+
+            width = values['width']
+            calculated_ar = float(width) / float(height)
+
+            aspect_ratio = Match(match.start, match.end, input_string=match.input_string,
+                                 name='aspect_ratio', value=round(calculated_ar, 3))
+            to_append.append(aspect_ratio)
+            if height in self.standard_heights and self.min_ar < calculated_ar < self.max_ar:
+                match.value = '{0}{1}'.format(height, scan_type)
+            else:
+                match.value = '{0}x{1}'.format(width, height)
+
+        return to_append
+
+
 class ScreenSizeOnlyOne(Rule):
     """
     Keep a single screen_size per filepath part.
diff --git a/guessit/test/rules/screen_size.yml b/guessit/test/rules/screen_size.yml
index 91db751..6260c81 100644
--- a/guessit/test/rules/screen_size.yml
+++ b/guessit/test/rules/screen_size.yml
@@ -4,11 +4,24 @@
 ? +360px
 ? "+360"
 ? +500x360
+? -250x360
 : screen_size: 360p
 
+? +640x360
+? -640x360i
+? -684x360i
+: screen_size: 360p
+  aspect_ratio: 1.778
+
 ? +360i
 : screen_size: 360i
 
+? +480x360i
+? -480x360p
+? -450x360
+: screen_size: 360i
+  aspect_ratio: 1.333
+
 ? +368p
 ? +368px
 ? -368i
@@ -16,23 +29,84 @@
 ? +500x368
 : screen_size: 368p
 
+? -490x368
+? -700x368
+: screen_size: 368p
+
+? +492x368p
+: screen_size:
+  aspect_ratio: 1.337
+
+? +654x368
+: screen_size: 368p
+  aspect_ratio: 1.777
+
+? +698x368
+: screen_size: 368p
+  aspect_ratio: 1.897
+
+? -368i
+? -368
+: screen_size: 368i
+
 ? +480p
 ? +480px
 ? -480i
 ? "+480"
-? +500x480
+? -500x480
+? -638x480
+? -920x480
 : screen_size: 480p
 
+? +640x480
+: screen_size: 480p
+  aspect_ratio: 1.333
+
+? +852x480
+: screen_size: 480p
+  aspect_ratio: 1.775
+
+? +910x480
+: screen_size: 480p
+  aspect_ratio: 1.896
+
+? +500x480
+? +500 x 480
+? +500 * 480
+? +500x480p
+? +500X480i
+: screen_size: 500x480
+  aspect_ratio: 1.042
+
 ? +480i
+? +852x480i
 : screen_size: 480i
 
 ? +576p
 ? +576px
 ? -576i
 ? "+576"
-? +500x576
+? -500x576
+? -766x576
+? -1094x576
 : screen_size: 576p
 
+? +768x576
+: screen_size: 576p
+  aspect_ratio: 1.333
+
+? +1024x576
+: screen_size: 576p
+  aspect_ratio: 1.778
+
+? +1092x576
+: screen_size: 576p
+  aspect_ratio: 1.896
+
+? +500x576
+: screen_size: 500x576
+  aspect_ratio: 0.868
+
 ? +576i
 : screen_size: 576i
 
@@ -42,16 +116,54 @@
 ? 720hd
 ? 720pHD
 ? "+720"
-? +500x720
+? -500x720
+? -950x720
+? -1368x720
 : screen_size: 720p
 
+? +960x720
+: screen_size: 720p
+  aspect_ratio: 1.333
+
+? +1280x720
+: screen_size: 720p
+  aspect_ratio: 1.778
+
+? +1366x720
+: screen_size: 720p
+  aspect_ratio: 1.897
+
+? +500x720
+: screen_size: 500x720
+  aspect_ratio: 0.694
+
 ? +900p
 ? +900px
 ? -900i
 ? "+900"
-? +500x900
+? -500x900
+? -1198x900
+? -1710x900
 : screen_size: 900p
 
+? +1200x900
+: screen_size: 900p
+  aspect_ratio: 1.333
+
+? +1600x900
+: screen_size: 900p
+  aspect_ratio: 1.778
+
+? +1708x900
+: screen_size: 900p
+  aspect_ratio: 1.898
+
+? +500x900
+? +500x900p
+? +500x900i
+: screen_size: 500x900
+  aspect_ratio: 0.556
+
 ? +900i
 : screen_size: 900i
 
@@ -61,27 +173,85 @@
 ? +1080pHD
 ? -1080i
 ? "+1080"
-? +500x1080
+? -500x1080
+? -1438x1080
+? -2050x1080
 : screen_size: 1080p
 
+? +1440x1080
+: screen_size: 1080p
+  aspect_ratio: 1.333
+
+? +1920x1080
+: screen_size: 1080p
+  aspect_ratio: 1.778
+
+? +2048x1080
+: screen_size: 1080p
+  aspect_ratio: 1.896
+
 ? +1080i
 ? -1080p
 : screen_size: 1080i
 
+? +500x1080
+: screen_size: 500x1080
+  aspect_ratio: 0.463
+
 ? +2160p
 ? +2160px
 ? -2160i
 ? "+2160"
 ? +4096x2160
+? +4k
+? -2878x2160
+? -4100x2160
 : screen_size: 2160p
 
+? +2880x2160
+: screen_size: 2160p
+  aspect_ratio: 1.333
+
+? +3840x2160
+: screen_size: 2160p
+  aspect_ratio: 1.778
+
+? +4098x2160
+: screen_size: 2160p
+  aspect_ratio: 1.897
+
+? +500x2160
+: screen_size: 500x2160
+  aspect_ratio: 0.231
+
 ? +4320p
 ? +4320px
 ? -4320i
 ? "+4320"
-? +7680x4320
+? -5758x2160
+? -8198x2160
 : screen_size: 4320p
 
+? +5760x4320
+: screen_size: 4320p
+  aspect_ratio: 1.333
+
+? +7680x4320
+: screen_size: 4320p
+  aspect_ratio: 1.778
+
+? +8196x4320
+: screen_size: 4320p
+  aspect_ratio: 1.897
+
+? +500x4320
+: screen_size: 500x4320
+  aspect_ratio: 0.116
+
 ? Test.File.720hd.bluray
+? Test.File.720p24
+? Test.File.720p30
 ? Test.File.720p50
+? Test.File.720p60
+? Test.File.720p120
 : screen_size: 720p