From 9cdc8325ae8ace6fe3b09eb78751d6912410e9cc Mon Sep 17 00:00:00 2001 From: Ian Cordasco Date: Sun, 10 Feb 2013 17:43:37 -0500 Subject: [PATCH 1/3] Fix Setting a cookie on redirect --- requests/sessions.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/requests/sessions.py b/requests/sessions.py index 406d78eb..37b5e205 100644 --- a/requests/sessions.py +++ b/requests/sessions.py @@ -136,6 +136,9 @@ class SessionRedirectMixin(object): except KeyError: pass + if response.headers.get('Set-Cookie'): + prepared_request.headers['Cookie'] = response.headers.get('Set-Cookie') + resp = self.send( prepared_request, stream=stream, From 0fb13e0b6c092e1acf96e24cf01a621ea4ef5eac Mon Sep 17 00:00:00 2001 From: Ian Cordasco Date: Sun, 10 Feb 2013 17:49:49 -0500 Subject: [PATCH 2/3] And tests --- test_requests.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/test_requests.py b/test_requests.py index 0355a4ad..e1317f53 100644 --- a/test_requests.py +++ b/test_requests.py @@ -111,6 +111,11 @@ class RequestsTestCase(unittest.TestCase): r = requests.get(httpbin('get') + '?test=true', params={'q': 'test'}, headers=heads) self.assertEqual(r.status_code, 200) + def test_set_cookie_on_301(self): + url = httpbin('cookies/set/foo/bar') + r = s.get(url) + self.assertTrue(s.cookies['foo'] == 'bar') + def test_user_agent_transfers(self): heads = { From 2e3169615636ceffc7aa85469464470c11d3a85e Mon Sep 17 00:00:00 2001 From: Ian Cordasco Date: Sun, 10 Feb 2013 19:36:36 -0500 Subject: [PATCH 3/3] Test and perfection for cookie handling. I also fixed up some of the RequestsCookieJar methods so using jar.update(other_jar) works without a problem. This cleans up some of the code in sessions and the resolve_redirects method. --- requests/cookies.py | 6 ++++++ requests/sessions.py | 11 ++++++++--- test_requests.py | 3 ++- 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/requests/cookies.py b/requests/cookies.py index bd7289ee..eb6c3146 100644 --- a/requests/cookies.py +++ b/requests/cookies.py @@ -240,12 +240,18 @@ class RequestsCookieJar(cookielib.CookieJar, collections.MutableMapping): """Dict-like __getitem__() for compatibility with client code. Throws exception if there are more than one cookie with name. In that case, use the more explicit get() method instead. Caution: operation is O(n), not O(1).""" + if isinstance(name, cookielib.Cookie): + name = name.name + return self._find_no_duplicates(name) def __setitem__(self, name, value): """Dict-like __setitem__ for compatibility with client code. Throws exception if there is already a cookie of that name in the jar. In that case, use the more explicit set() method instead.""" + if isinstance(name, cookielib.Cookie): + name = name.name + self.set(name, value) def __delitem__(self, name): diff --git a/requests/sessions.py b/requests/sessions.py index 37b5e205..ff70db74 100644 --- a/requests/sessions.py +++ b/requests/sessions.py @@ -85,6 +85,7 @@ class SessionRedirectMixin(object): prepared_request.hooks = req.hooks prepared_request.method = req.method prepared_request.url = req.url + cookiejar = resp.cookies # ((resp.status_code is codes.see_other)) while (('location' in resp.headers and resp.status_code in REDIRECT_STATI)): @@ -131,13 +132,13 @@ class SessionRedirectMixin(object): prepared_request.body = None + headers = prepared_request.headers try: - del prepared_request.headers['Cookie'] + del headers['Cookie'] except KeyError: pass - if response.headers.get('Set-Cookie'): - prepared_request.headers['Cookie'] = response.headers.get('Set-Cookie') + prepared_request.prepare_cookies(cookiejar) resp = self.send( prepared_request, @@ -149,9 +150,13 @@ class SessionRedirectMixin(object): allow_redirects=False, ) + cookiejar.update(resp.cookies) + i += 1 yield resp + resp.cookies.update(cookiejar) + class Session(SessionRedirectMixin): """A Requests session. diff --git a/test_requests.py b/test_requests.py index e1317f53..6c4c3dd2 100644 --- a/test_requests.py +++ b/test_requests.py @@ -112,7 +112,8 @@ class RequestsTestCase(unittest.TestCase): self.assertEqual(r.status_code, 200) def test_set_cookie_on_301(self): - url = httpbin('cookies/set/foo/bar') + s = requests.session() + url = httpbin('cookies/set?foo=bar') r = s.get(url) self.assertTrue(s.cookies['foo'] == 'bar')