Add 3d support for searching

couchpotato/core/media/_base/searcher/main.py

@@ -17,6 +17,7 @@ class Searcher(SearcherBase):
     def __init__(self):
         addEvent('searcher.protocols', self.getSearchProtocols)
         addEvent('searcher.contains_other_quality', self.containsOtherQuality)
+        addEvent('searcher.correct_3d', self.correct3D)
         addEvent('searcher.correct_year', self.correctYear)
         addEvent('searcher.correct_name', self.correctName)
         addEvent('searcher.correct_words', self.correctWords)
@@ -123,6 +124,17 @@ class Searcher(SearcherBase):
 
         return not (found.get(preferred_quality['identifier']) and len(found) == 1)
 
+    def correct3D(self, nzb, preferred_quality = None):
+        if not preferred_quality: preferred_quality = {}
+        if not preferred_quality.get('custom'): return
+
+        threed = preferred_quality['custom'].get('3d')
+
+        # Try guessing via quality tags
+        guess = fireEvent('quality.guess', [nzb.get('name')], single = True)
+
+        return threed == guess.get('is_3d')
+
     def correctYear(self, haystack, year, year_range):
 
         if not isinstance(haystack, (list, tuple, set)):

couchpotato/core/media/movie/_base/static/movie.actions.js

@@ -192,7 +192,7 @@ MA.Release = new Class({
 					}).adopt(
 						new Element('span.name', {'text': release_name, 'title': release_name}),
 						new Element('span.status', {'text': release.status, 'class': 'release_status '+release.status}),
-						new Element('span.quality', {'text': quality.label || 'n/a'}),
+						new Element('span.quality', {'text': quality.label + (release.is_3d ? ' 3D' : '') || 'n/a'}),
 						new Element('span.size', {'text': info['size'] ? Math.floor(self.get(release, 'size')) : 'n/a'}),
 						new Element('span.age', {'text': self.get(release, 'age')}),
 						new Element('span.score', {'text': self.get(release, 'score')}),

couchpotato/core/media/movie/searcher.py

@@ -143,14 +143,17 @@ class MovieSearcher(SearcherBase, MovieTypeBase):
 
         ret = False
 
+        index = 0
         for q_identifier in profile.get('qualities'):
-            index = profile['qualities'].index(q_identifier)
             quality_custom = {
                 'quality': q_identifier,
                 'finish': profile['finish'][index],
-                'wait_for': profile['wait_for'][index]
+                'wait_for': profile['wait_for'][index],
+                '3d': profile['3d'][index] if profile['3d'] else False
             }
 
+            index += 1
+
             if not self.conf('always_search') and not self.couldBeReleased(q_identifier in pre_releases, release_dates, movie['info']['year']):
                 too_early_to_search.append(q_identifier)
                 continue
@@ -168,6 +171,9 @@ class MovieSearcher(SearcherBase, MovieTypeBase):
                 quality = fireEvent('quality.single', identifier = q_identifier, single = True)
                 log.info('Search for %s in %s', (default_title, quality['label']))
 
+                # Extend quality with profile customs
+                quality['custom'] = quality_custom
+
                 results = fireEvent('searcher.search', search_protocols, movie, quality, single = True) or []
                 if len(results) == 0:
                     log.debug('Nothing found for %s in %s', (default_title, quality['label']))
@@ -221,13 +227,17 @@ class MovieSearcher(SearcherBase, MovieTypeBase):
         if not fireEvent('searcher.correct_words', nzb['name'], media, single = True):
             return False
 
-        preferred_quality = fireEvent('quality.single', identifier = quality['identifier'], single = True)
+        preferred_quality = quality if quality else fireEvent('quality.single', identifier = quality['identifier'], single = True)
 
         # Contains lower quality string
         if fireEvent('searcher.contains_other_quality', nzb, movie_year = media['info']['year'], preferred_quality = preferred_quality, single = True):
             log.info2('Wrong: %s, looking for %s', (nzb['name'], quality['label']))
             return False
 
+        # Contains lower quality string
+        if not fireEvent('searcher.correct_3d', nzb, preferred_quality = preferred_quality, single = True):
+            log.info2('Wrong: %s, %slooking for %s in 3D', (nzb['name'], ('' if preferred_quality['custom'].get('3d') else 'NOT '), quality['label']))
+            return False
 
         # File to small
         if nzb['size'] and tryInt(preferred_quality['size_min']) > tryInt(nzb['size']):

couchpotato/core/plugins/profile/main.py

@@ -188,11 +188,11 @@ class ProfilePlugin(Plugin):
                 'qualities': ['dvdrip', 'dvdr']
             }, {
                 'label': 'Prefer 3D HD',
-                'qualities': ['720p', '1080p', '720p', '1080p'],
+                'qualities': ['1080p', '720p', '720p', '1080p'],
                 '3d': [True, True]
             }, {
                 'label': '3D HD',
-                'qualities': ['720p', '1080p'],
+                'qualities': ['1080p', '720p'],
                 '3d': [True, True]
             }]
 

couchpotato/core/plugins/quality/main.py

@@ -190,14 +190,19 @@ class QualityPlugin(Plugin):
         # Start with 0
         score = {}
         for quality in qualities:
-            score[quality.get('identifier')] = 0
+            score[quality.get('identifier')] = {
+                'score': 0,
+                '3d': {}
+            }
 
         for cur_file in files:
             words = re.split('\W+', cur_file.lower())
 
             for quality in qualities:
                 contains_score = self.containsTagScore(quality, words, cur_file)
-                self.calcScore(score, quality, contains_score)
+                threedscore = self.contains3D(quality, words, cur_file) if quality.get('allow_3d') else (0, None)
+
+                self.calcScore(score, quality, contains_score, threedscore)
 
         # Try again with loose testing
         for quality in qualities:
@@ -213,10 +218,13 @@ class QualityPlugin(Plugin):
         if not has_non_zero:
             return None
 
-        heighest_quality = max(score, key = score.get)
+        heighest_quality = max(score, key = lambda p: score[p]['score'])
         if heighest_quality:
             for quality in qualities:
                 if quality.get('identifier') == heighest_quality:
+                    quality['is_3d'] = False
+                    if score[heighest_quality].get('3d'):
+                        quality['is_3d'] = True
                     return self.setCache(cache_key, quality)
 
         return None
@@ -260,6 +268,23 @@ class QualityPlugin(Plugin):
 
         return score
 
+    def contains3D(self, quality, words, cur_file = ''):
+        cur_file = ss(cur_file)
+
+        for key in self.threed_tags:
+            tags = self.threed_tags.get(key, [])
+
+            for tag in tags:
+                if (isinstance(tag, tuple) and '.'.join(tag) in '.'.join(words)) or (isinstance(tag, (str, unicode)) and ss(tag.lower()) in cur_file.lower()):
+                    log.debug('Found %s in %s', (tag, cur_file))
+                    return 1, key
+
+            if list(set([key]) & set(words)):
+                log.debug('Found %s in %s', (tag, cur_file))
+                return 1, key
+
+        return 0, None
+
     def guessLooseScore(self, quality, extra = None):
 
         score = 0
@@ -282,9 +307,16 @@ class QualityPlugin(Plugin):
 
         return score
 
-    def calcScore(self, score, quality, add_score):
+    def calcScore(self, score, quality, add_score, threedscore = (0, None)):
+
+        score[quality['identifier']]['score'] += add_score
+
+        threedscore, threedtag = threedscore
+        if threedscore and threedtag:
+            if threedscore not in score[quality['identifier']]['3d']:
+                score[quality['identifier']]['3d'][threedtag] = 0
 
-        score[quality['identifier']] += add_score
+            score[quality['identifier']]['3d'][threedtag] += threedscore
 
         # Set order for allow calculation (and cache)
         if not self.cached_order:
@@ -294,7 +326,7 @@ class QualityPlugin(Plugin):
 
         if add_score != 0:
             for allow in quality.get('allow', []):
-                score[allow] -= 40 if self.cached_order[allow] < self.cached_order[quality['identifier']] else 5
+                score[allow]['score'] -= 40 if self.cached_order[allow] < self.cached_order[quality['identifier']] else 5
 
     def doTest(self):
 

couchpotato/core/plugins/release/main.py

@@ -343,6 +343,10 @@ class Release(Plugin):
 
             found_releases = []
 
+            is_3d = False
+            try: is_3d = quality['custom']['3d']
+            except: pass
+
             for rel in search_results:
 
                 rel_identifier = md5(rel['url'])
@@ -353,6 +357,7 @@ class Release(Plugin):
                     'identifier': rel_identifier,
                     'media_id': media.get('_id'),
                     'quality': quality.get('identifier'),
+                    'is_3d': is_3d,
                     'status': rel.get('status', 'available'),
                     'last_edit': int(time.time()),
                     'info': {}