diff --git a/couchpotato/core/providers/movie/_modifier/main.py b/couchpotato/core/providers/movie/_modifier/main.py index 4f58f400..d62f8fe0 100644 --- a/couchpotato/core/providers/movie/_modifier/main.py +++ b/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]['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()) + temp[imdb] = self.getLibraryTags(imdb) # 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']))