Merge pull request #1180 from sigmavirus24/master

Fix Session level Cookie Handling

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):

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,11 +132,14 @@ class SessionRedirectMixin(object):
 
                 prepared_request.body = None
 
+            headers = prepared_request.headers
             try:
-                del prepared_request.headers['Cookie']
+                del headers['Cookie']
             except KeyError:
                 pass
 
+            prepared_request.prepare_cookies(cookiejar)
+
             resp = self.send(
                 prepared_request,
                 stream=stream,
@@ -146,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.

test_requests.py

@@ -111,6 +111,12 @@ 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):
+        s = requests.session()
+        url = httpbin('cookies/set?foo=bar')
+        r = s.get(url)
+        self.assertTrue(s.cookies['foo'] == 'bar')
+
     def test_user_agent_transfers(self):
 
         heads = {