diff --git a/couchpotato/core/media/__init__.py b/couchpotato/core/media/__init__.py index 1ba8386..84545e4 100644 --- a/couchpotato/core/media/__init__.py +++ b/couchpotato/core/media/__init__.py @@ -11,7 +11,7 @@ class MediaBase(Plugin): default_dict = { 'profile': {'types': {'quality': {}}}, 'releases': {'status': {}, 'quality': {}, 'files':{}, 'info': {}}, - 'library': {'titles': {}, 'files':{}}, + 'library': {'titles': {}, 'files':{}, 'related_libraries': {}}, 'files': {}, 'status': {}, 'category': {}, diff --git a/couchpotato/core/media/show/searcher/main.py b/couchpotato/core/media/show/searcher/main.py index 1df6c6f..33e5aff 100644 --- a/couchpotato/core/media/show/searcher/main.py +++ b/couchpotato/core/media/show/searcher/main.py @@ -40,7 +40,6 @@ class ShowSearcher(Plugin): addEvent('searcher.correct_release', self.correctRelease) addEvent('searcher.get_media_identifier', self.getMediaIdentifier) - addEvent('searcher.get_media_root', self.getMediaRoot) def single(self, media, search_protocols = None, manual = False): if media['type'] == 'show': @@ -73,6 +72,7 @@ class ShowSearcher(Plugin): #fireEvent('episode.delete', episode['id'], single = True) return + # TODO replace with 'related_libraries' show, season, episode = self.getMedia(media) if show is None or season is None: log.error('Unable to find show or season library in database, missing required data for searching') @@ -139,6 +139,7 @@ class ShowSearcher(Plugin): if media['type'] not in ['show', 'season', 'episode']: return + # TODO replace with 'related_libraries' show, season, episode = self.getMedia(media) if show is None: return None @@ -199,6 +200,7 @@ class ShowSearcher(Plugin): if not fireEvent('searcher.correct_words', release['name'], media, single = True): return False + # TODO replace with 'related_libraries' show, season, episode = self.getMedia(media) if show is None or season is None: log.error('Unable to find show or season library in database, missing required data for searching') @@ -259,18 +261,6 @@ class ShowSearcher(Plugin): return identifier # TODO move this somewhere else - def getMediaRoot(self, media): - if media['type'] not in ['show', 'season', 'episode']: - return None - - show, season, episode = self.getMedia(media) - if show is None or season is None: - log.error('Unable to find show or season library in database, missing required data for searching') - return - - return show.to_dict() - - # TODO move this somewhere else def getMedia(self, media): db = get_session() diff --git a/couchpotato/core/plugins/matcher/main.py b/couchpotato/core/plugins/matcher/main.py index c20db2b..ccd46d3 100644 --- a/couchpotato/core/plugins/matcher/main.py +++ b/couchpotato/core/plugins/matcher/main.py @@ -73,6 +73,7 @@ class Matcher(Plugin): return True def correctTitle(self, chain, media): + # TODO need to switch this to use media['related_libraries'] root_library = fireEvent('searcher.get_media_root', media['library'], single = True) if 'show_name' not in chain.info or not len(chain.info['show_name']): diff --git a/couchpotato/core/settings/model.py b/couchpotato/core/settings/model.py index 472f669..9f13f8a 100644 --- a/couchpotato/core/settings/model.py +++ b/couchpotato/core/settings/model.py @@ -112,6 +112,40 @@ class Library(Entity): parent = ManyToOne('Library') children = OneToMany('Library') + def getRelated(self, include_parents = True, include_self = True, include_children = True): + libraries = [] + + if include_parents and self.parent is not None: + libraries += self.parent.getRelated(include_children = False) + + if include_self: + libraries += [(self.type, self)] + + if include_children: + for child in self.children: + libraries += child.getRelated(include_parents = False) + + return libraries + + def to_dict(self, deep = None, exclude = None): + if not exclude: exclude = [] + if not deep: deep = {} + + include_related = deep.pop('related_libraries', None) is not None + + orig_dict = super(Library, self).to_dict(deep = deep, exclude = exclude) + + # Include related libraries (parents and children) + if include_related and self.parent is not None: + # TODO need a way to flag the root library + # TODO converting to a dict will remove libraries with multiple types (episodes, seasons), they need to be merged into a list instead + orig_dict['related_libraries'] = dict([ + (library_type, library.to_dict(deep, exclude)) + for (library_type, library) in self.getRelated(include_self = False) + ]) + + return orig_dict + class ShowLibrary(Library, DictMixin): using_options(inheritance = 'multi')