diff --git a/weboob/tools/browser2/sessions.py b/weboob/tools/browser2/sessions.py index de3dc25dea0db8aa9c14d06019c7a7683a81721e..3ee4757f09f6820db57dfc3ec06f5401f4b02219 100644 --- a/weboob/tools/browser2/sessions.py +++ b/weboob/tools/browser2/sessions.py @@ -27,8 +27,79 @@ from requests import Session from requests.adapters import DEFAULT_POOLSIZE, HTTPAdapter +from requests.compat import cookielib, OrderedDict +from requests.cookies import cookiejar_from_dict, RequestsCookieJar +from requests.models import PreparedRequest +from requests.sessions import merge_setting +from requests.structures import CaseInsensitiveDict +from requests.utils import get_netrc_auth -class FuturesSession(Session): +def merge_hooks(request_hooks, session_hooks, dict_class=OrderedDict): + """ + Properly merges both requests and session hooks. + + This is necessary because when request_hooks == {'response': []}, the + merge breaks Session hooks entirely. + + Backport from request so we can use it in wheezy + """ + if session_hooks is None or session_hooks.get('response') == []: + return request_hooks + + if request_hooks is None or request_hooks.get('response') == []: + return session_hooks + + ret = {} + for (k, v) in request_hooks.items(): + if v is not None: + ret[k] = set(v).union(session_hooks.get(k, [])) + + return ret + + +class WeboobSession(Session): + + def prepare_request(self, request): + """Constructs a :class:`PreparedRequest ` for + transmission and returns it. The :class:`PreparedRequest` has settings + merged from the :class:`Request ` instance and those of the + :class:`Session`. + + :param request: :class:`Request` instance to prepare with this + session's settings. + """ + cookies = request.cookies or {} + + # Bootstrap CookieJar. + if not isinstance(cookies, cookielib.CookieJar): + cookies = cookiejar_from_dict(cookies) + + # Merge with session cookies + merged_cookies = RequestsCookieJar() + merged_cookies.update(self.cookies) + merged_cookies.update(cookies) + + + # Set environment's basic authentication if not explicitly set. + auth = request.auth + if self.trust_env and not auth and not self.auth: + auth = get_netrc_auth(request.url) + + p = PreparedRequest() + p.prepare( + method=request.method.upper(), + url=request.url, + files=request.files, + data=request.data, + headers=merge_setting(request.headers, self.headers, dict_class=CaseInsensitiveDict), + params=merge_setting(request.params, self.params), + auth=merge_setting(auth, self.auth), + cookies=merged_cookies, + hooks=merge_hooks(request.hooks, self.hooks), + ) + return p + +class FuturesSession(WeboobSession): def __init__(self, executor=None, max_workers=2, *args, **kwargs): """Creates a FuturesSession