1# -*- coding: utf-8 -*-
2
3# Copyright(C) 2017  Vincent A
4#
5# This file is part of weboob.
6#
7# weboob is free software: you can redistribute it and/or modify
8# it under the terms of the GNU Lesser General Public License as published by
9# the Free Software Foundation, either version 3 of the License, or
10# (at your option) any later version.
11#
12# weboob is distributed in the hope that it will be useful,
13# but WITHOUT ANY WARRANTY; without even the implied warranty of
14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15# GNU Lesser General Public License for more details.
16#
17# You should have received a copy of the GNU Lesser General Public License
18# along with weboob. If not, see <http://www.gnu.org/licenses/>.
19
20from contextlib import contextmanager
21from functools import wraps
22
23from .browsers import LoginBrowser
24from .exceptions import LoggedOut
25from ..exceptions import BrowserUnavailable
26
27
28__all__ = ['login_method', 'retry_on_logout', 'RetryLoginBrowser']
29
30
31def login_method(func):
32    """Decorate a method to indicate the browser is logging in.
33
34    When the decorated method is called, pages like
35    `weboob.browser.pages.LoginPage` will not raise `LoggedOut`, since it is
36    expected for the browser not to be logged yet.
37    """
38
39    func.login_decorated = True
40
41    @wraps(func)
42    def wrapper(browser, *args, **kwargs):
43        browser.logging_in += 1
44        try:
45            return func(browser, *args, **kwargs)
46        finally:
47            browser.logging_in -= 1
48    return wrapper
49
50
51def retry_on_logout(exc_check=LoggedOut, tries=4):
52    """Decorate a function to retry several times in case of exception.
53
54    The decorated function is called at max 4 times. It is retried only when it
55    raises an exception of the type `weboob.browser.exceptions.LoggedOut`.
56    If the function call succeeds and returns an iterator, a wrapper to the
57    iterator is returned. If iterating on the result raises a `LoggedOut`,
58    the iterator is recreated by re-calling the function, but the values
59    already yielded will not be re-yielded.
60    For consistency, the function MUST always return values in the same order.
61
62    Adding this decorator to a method which can be called from another
63    decorated method should be avoided, since nested calls will greatly
64    increase the number of retries.
65    """
66
67    if not isinstance(exc_check, type) or not issubclass(exc_check, Exception):
68        raise TypeError('retry_on_logout() must be called in order to decorate %r' % tries)
69
70    def decorator(func):
71        @wraps(func)
72        def wrapper(browser, *args, **kwargs):
73            cb = lambda: func(browser, *args, **kwargs)
74
75            for i in range(tries, 0, -1):
76                try:
77                    ret = cb()
78                except exc_check as exc:
79                    browser.logger.info('%r raised, retrying', exc)
80                    continue
81
82                if not (hasattr(ret, '__next__') or hasattr(ret, 'next')):
83                    return ret  # simple value, no need to retry on items
84                return iter_retry(cb, value=ret, remaining=i, exc_check=exc_check, logger=browser.logger)
85
86            raise BrowserUnavailable('Site did not reply successfully after multiple tries')
87
88        return wrapper
89    return decorator
90
91
92@contextmanager
93def retry_on_logout_context(tries=4, logger=None):
94    for i in range(tries, 0, -1):
95        try:
96            yield
97        except LoggedOut as exc:
98            if logger:
99                logger.debug('%r raised, retrying', exc)
100        else:
101            return
102    raise BrowserUnavailable('Site did not reply successfully after multiple tries')
103
104
105class RetryLoginBrowser(LoginBrowser):
106    """Browser that can retry methods if the site logs out the session.
107
108    Some sites can terminate a session anytime, redirecting to a login page.
109    To avoid having to handle it in the middle of every method, one can simply
110    let logouts raise a `weboob.browser.exceptions.LoggedOut` exception that
111    is handled with a retry, thanks to the `@retry_on_logout` decorator.
112
113    The `weboob.browser.pages.LoginPage` will raise `LoggedOut` if the browser
114    is not currently logging in. To detect this situation, the `do_login`
115    method MUST be decorated with `@login_method`.
116    """
117    def __init__(self, *args, **kwargs):
118        super(RetryLoginBrowser, self).__init__(*args, **kwargs)
119        self.logging_in = 0
120
121        if not hasattr(self.do_login, 'login_decorated'):
122            raise Exception('do_login method was not decorated with @login_method')
123
124
125class iter_retry(object):
126    # when the callback is retried, it will create a new iterator, but we may already yielded
127    # some values, so we need to keep track of them and seek in the middle of the iterator
128
129    def __init__(self, cb, remaining=4, value=None, exc_check=LoggedOut, logger=None):
130        self.cb = cb
131        self.it = iter(value) if value is not None else None
132        self.items = []
133        self.remaining = remaining
134        self.logger = logger
135        self.exc_check = exc_check
136
137    def __iter__(self):
138        return self
139
140    def __next__(self):
141        if self.remaining <= 0:
142            raise BrowserUnavailable('Site did not reply successfully after multiple tries')
143
144        if self.it is None:
145            self.it = iter(self.cb())
146
147            # recreated iterator, consume previous items
148            try:
149                nb = -1
150                for nb, sent in enumerate(self.items):
151                    new = next(self.it)
152                    if hasattr(new, 'iter_fields'):
153                        equal = dict(sent.iter_fields()) == dict(new.iter_fields())
154                    else:
155                        equal = sent == new
156                    if not equal:
157                        # safety is not guaranteed
158                        raise BrowserUnavailable('Site replied inconsistently between retries, %r vs %r', sent, new)
159            except StopIteration:
160                raise BrowserUnavailable('Site replied fewer elements (%d) than last iteration (%d)', nb + 1, len(self.items))
161            except self.exc_check as exc:
162                if self.logger:
163                    self.logger.info('%r raised, retrying', exc)
164                self.it = None
165                self.remaining -= 1
166                return next(self)
167
168        # return one item
169        try:
170            obj = next(self.it)
171        except self.exc_check as exc:
172            if self.logger:
173                self.logger.info('%r raised, retrying', exc)
174            self.it = None
175            self.remaining -= 1
176            return next(self)
177        else:
178            self.items.append(obj)
179            return obj
180
181    next = __next__
182