diff --git a/weboob/backends/aum/backend.py b/weboob/backends/aum/backend.py index 27302375..5707e725 100644 --- a/weboob/backends/aum/backend.py +++ b/weboob/backends/aum/backend.py @@ -140,7 +140,7 @@ class AuMBackend(BaseBackend, ICapMessages, ICapMessagesReply, ICapDating, ICapC new_baskets = self.browser.nb_new_baskets() if new_baskets: ids = self.browser.get_baskets() - while new_baskets > 0: + while new_baskets > 0 and len(ids) > new_baskets: new_baskets -= 1 profile = self.browser.get_profile(ids[new_baskets]) diff --git a/weboob/backends/youjizz/backend.py b/weboob/backends/youjizz/backend.py index 75f23ed7..3f143e09 100644 --- a/weboob/backends/youjizz/backend.py +++ b/weboob/backends/youjizz/backend.py @@ -37,19 +37,23 @@ class YoujizzBackend(BaseBackend, ICapVideo): BROWSER = YoujizzBrowser def get_video(self, _id): - video = self.browser.get_video(_id) + with self.browser: + video = self.browser.get_video(_id) return video def iter_page_urls(self, mozaic_url): - return self.browser.iter_page_urls(mozaic_url) + with self.browser: + return self.browser.iter_page_urls(mozaic_url) def iter_search_results(self, pattern=None, sortby=ICapVideo.SEARCH_RELEVANCE, nsfw=False): if not nsfw: return set() - return self.browser.iter_search_results(pattern) + with self.browser: + return self.browser.iter_search_results(pattern) def fill_video(self, video, fields): # ignore the fields param: VideoPage.get_video() returns all the information - return self.browser.get_video(YoujizzVideo.id2url(video.id), video) + with self.browser: + return self.browser.get_video(YoujizzVideo.id2url(video.id), video) OBJECTS = {YoujizzVideo: fill_video} diff --git a/weboob/backends/youporn/backend.py b/weboob/backends/youporn/backend.py index 5fe2262e..5f35c8d1 100644 --- a/weboob/backends/youporn/backend.py +++ b/weboob/backends/youporn/backend.py @@ -37,19 +37,22 @@ class YoupornBackend(BaseBackend, ICapVideo): BROWSER = YoupornBrowser def get_video(self, _id): - return self.browser.get_video(_id) + with self.browser: + return self.browser.get_video(_id) SORTBY = ['relevance', 'rating', 'views', 'time'] def iter_search_results(self, pattern=None, sortby=ICapVideo.SEARCH_RELEVANCE, nsfw=False): if not nsfw: return set() - return self.browser.iter_search_results(pattern, self.SORTBY[sortby]) + with self.browser: + return self.browser.iter_search_results(pattern, self.SORTBY[sortby]) def iter_page_urls(self, mozaic_url): raise NotImplementedError() def fill_video(self, video, fields): # ignore the fields param: VideoPage.get_video() returns all the information - return self.browser.get_video(YoupornVideo.id2url(video.id), video) + with self.browser: + return self.browser.get_video(YoupornVideo.id2url(video.id), video) OBJECTS = {YoupornVideo: fill_video} diff --git a/weboob/backends/youtube/backend.py b/weboob/backends/youtube/backend.py index ce681648..02626fa8 100644 --- a/weboob/backends/youtube/backend.py +++ b/weboob/backends/youtube/backend.py @@ -39,7 +39,8 @@ class YoutubeBackend(BaseBackend, ICapVideo): BROWSER = YoutubeBrowser def get_video(self, _id): - return self.browser.get_video(_id) + with self.browser: + return self.browser.get_video(_id) def iter_search_results(self, pattern=None, sortby=ICapVideo.SEARCH_RELEVANCE, nsfw=False): import gdata.youtube.service @@ -68,6 +69,7 @@ class YoutubeBackend(BaseBackend, ICapVideo): def fill_video(self, video, fields): # ignore the fields param: VideoPage.get_video() returns all the information - return self.browser.get_video(YoutubeVideo.id2url(video.id), video) + with self.browser: + return self.browser.get_video(YoutubeVideo.id2url(video.id), video) OBJECTS = {YoutubeVideo: fill_video} diff --git a/weboob/tools/application/base.py b/weboob/tools/application/base.py index 8478106e..8dc67ec4 100644 --- a/weboob/tools/application/base.py +++ b/weboob/tools/application/base.py @@ -23,6 +23,7 @@ from optparse import OptionGroup, OptionParser from weboob.core.ouiboube import Weboob from weboob.tools.config.iconfig import ConfigError +from weboob.tools.backend import ObjectNotSupported __all__ = ['BackendNotFound', 'BaseApplication', 'ConfigError'] @@ -185,6 +186,34 @@ class BaseApplication(object): version = '%s v%s' % (self.APPNAME, self.VERSION) return version + def _complete_obj(self, backend, obj, fields): + if fields: + try: + backend.fillobj(obj, fields) + except ObjectNotSupported, e: + logging.warning(u'Could not retrieve required fields (%s): %s' % (','.join(fields), e)) + return obj + + def _complete_iter(self, backend, count, fields, res): + for i, sub in enumerate(res): + if count and i == count: + break + sub = self._complete_obj(backend, sub, fields) + yield sub + + def complete(self, backend, count, selected_fields, function, *args, **kwargs): + res = getattr(backend, function)(*args, **kwargs) + + if self.selected_fields: + fields = set(self.selected_fields) - set('*') + else: + fields = None + + if hasattr(res, '__iter__'): + return self._complete_iter(backend, count, fields, res) + else: + return self._complete_obj(backend, fields, res) + @classmethod def run(klass, args=None): """ diff --git a/weboob/tools/application/console.py b/weboob/tools/application/console.py index eca5e0a2..3bb8293e 100644 --- a/weboob/tools/application/console.py +++ b/weboob/tools/application/console.py @@ -25,7 +25,6 @@ import sys from weboob.core import CallErrors from weboob.core.backends import BackendsConfig -from weboob.tools.backend import ObjectNotSupported from .base import BackendNotFound, BaseApplication from .formatters.load import formatters, load_formatter @@ -278,14 +277,5 @@ class ConsoleApplication(BaseApplication): """ Call Weboob.do(), after having filled the yielded object, if selected fields are given by user. """ - for i, (backend, result) in enumerate(self.weboob.do(function, *args, **kwargs)): - if self.options.count and i == self.options.count: - break - if self.selected_fields: - fields = set(self.selected_fields) - set('*') - if fields: - try: - backend.fillobj(result, fields) - except ObjectNotSupported, e: - logging.warning(u'Could not retrieve required fields (%s): %s' % (','.join(fields), e)) - yield backend, result + + return self.weboob.do(self.complete, self.options.count, self.selected_fields, function, *args, **kwargs)