Add "already in library/wanted" to userscript results

couchpotato/core/providers/movie/_modifier/main.py

@@ -13,40 +13,19 @@ class MovieResultModifier(Plugin):
 
     def __init__(self):
         addEvent('result.modify.movie.search', self.combineOnIMDB)
+        addEvent('result.modify.movie.info', self.checkLibrary)
 
     def combineOnIMDB(self, results):
 
         temp = {}
         unique = 1
 
-        # Statuses
-        active_status = fireEvent('status.get', 'active', single = True)
-        done_status = fireEvent('status.get', 'done', single = True)
-
         # Combine on imdb id
         for item in results:
             imdb = item.get('imdb')
             if imdb:
                 if not temp.get(imdb):
-                    temp[imdb] = {}
+                    temp[imdb] = self.getLibraryTags(imdb)
-
-                    temp[imdb]['in_wanted'] = False
-                    temp[imdb]['in_library'] = False
-
-                    # Add release info from current library
-                    try:
-                        db = get_session()
-                        l = db.query(Library).filter_by(identifier = imdb).first()
-                        if l:
-                            for movie in l.movies:
-                                if movie.status_id == active_status['id']:
-                                    temp[imdb]['in_wanted'] = movie.profile.to_dict()
-
-                                for release in movie.releases:
-                                    if release.status_id == done_status['id']:
-                                        temp[imdb]['in_library'] = release.quality.to_dict()
-                    except:
-                        log.error('Tried getting more info on searched movies: %s' % traceback.format_exc())
 
                 # Merge dicts
                 temp[imdb] = mergeDicts(temp[imdb], item)
@@ -58,3 +37,35 @@ class MovieResultModifier(Plugin):
         temp_list = [temp[x] for x in temp]
 
         return temp_list
+
+    def getLibraryTags(self, imdb):
+
+        temp = {
+            'in_wanted': False,
+            'in_library': False,
+        }
+
+        # Add release info from current library
+        try:
+            db = get_session()
+            l = db.query(Library).filter_by(identifier = imdb).first()
+            if l:
+
+                # Statuses
+                active_status = fireEvent('status.get', 'active', single = True)
+                done_status = fireEvent('status.get', 'done', single = True)
+
+                for movie in l.movies:
+                    if movie.status_id == active_status['id']:
+                        temp['in_wanted'] = movie.profile.to_dict()
+
+                    for release in movie.releases:
+                        if release.status_id == done_status['id']:
+                            temp['in_library'] = release.quality.to_dict()
+        except:
+            log.error('Tried getting more info on searched movies: %s' % traceback.format_exc())
+
+        return temp
+
+    def checkLibrary(self, result):
+        return mergeDicts(result, self.getLibraryTags(result['imdb']))