Merge pull request #626 from joshimhoff/develop

Improvements to RequestsCookieJar

requests/cookies.py

@@ -120,6 +120,10 @@ def remove_cookie_by_name(cookiejar, name, domain=None, path=None):
     for domain, path, name in clearables:
         cookiejar.clear(domain, path, name)
 
+class CookieConflictError(RuntimeError):
+    """There are two cookies that meet the criteria specified in the cookie jar. 
+    Use .get and .set and include domain and path args in order to be more specific."""
+
 class RequestsCookieJar(cookielib.CookieJar, collections.MutableMapping):
     """Compatibility class; is a cookielib.CookieJar, but exposes a dict interface.
 
@@ -138,12 +142,18 @@ class RequestsCookieJar(cookielib.CookieJar, collections.MutableMapping):
     """
 
     def get(self, name, default=None, domain=None, path=None):
+        """Dict-like get() that also supports optional domain and path args in
+        order to resolve naming collisions from using one cookie jar over
+        multiple domains. Caution: operation is O(n), not O(1)."""
         try:
-            return self._find(name, domain, path)
+            return self._find_no_duplicates(name, domain, path)
         except KeyError:
             return default
 
     def set(self, name, value, **kwargs):
+        """Dict-like set() that also supports optional domain and path args in
+        order to resolve naming collisions from using one cookie jar over
+        multiple domains."""
         # support client code that unsets cookies by assignment of a None value:
         if value is None:
             remove_cookie_by_name(self, name, domain=kwargs.get('domain'), path=kwargs.get('path'))
@@ -156,16 +166,88 @@ class RequestsCookieJar(cookielib.CookieJar, collections.MutableMapping):
         self.set_cookie(c)
         return c
 
+    def keys(self):
+        """Dict-like keys() that returns a list of names of cookies from the jar.
+        See values() and items()."""
+        keys = []
+        for cookie in iter(self):
+            keys.append(cookie.name)
+        return keys
+
+    def values(self):
+        """Dict-like values() that returns a list of values of cookies from the jar.
+        See keys() and items()."""
+        values = []
+        for cookie in iter(self):
+            values.append(cookie.value)
+        return values
+    
+    def items(self):
+        """Dict-like items() that returns a list of name-value tuples from the jar.
+        See keys() and values(). Allows client-code to call "dict(RequestsCookieJar)
+        and get a vanilla python dict of key value pairs."""
+        items = []
+        for cookie in iter(self):
+            items.append((cookie.name, cookie.value))
+        return items
+
+    def list_domains(self):
+        """Utility method to list all the domains in the jar."""
+        domains = []
+        for cookie in iter(self):
+            if cookie.domain not in domains:
+                domains.append(cookie.domain)
+        return domains
+
+    def list_paths(self):
+        """Utility method to list all the paths in the jar."""
+        paths = []
+        for cookie in iter(self):
+            if cookie.path not in paths:
+                paths.append(cookie.path)
+        return paths
+
+    def multiple_domains(self):
+        """Returns True if there are multiple domains in the jar.
+        Returns False otherwise."""
+        domains = []
+        for cookie in iter(self):
+            if cookie.domain is not None and cookie.domain in domains:
+                return True
+            domains.append(cookie.domain)
+        return False # there is only one domain in jar
+
+    def get_dict(self, domain=None, path=None):
+        """Takes as an argument an optional domain and path and returns a plain old
+        Python dict of name-value pairs of cookies that meet the requirements."""
+        dictionary = {}
+        for cookie in iter(self):
+            if (domain == None or cookie.domain == domain) and (path == None 
+                                                or cookie.path == path):
+                dictionary[cookie.name] = cookie.value
+        return dictionary
+
     def __getitem__(self, name):
-        return self._find(name)
+        """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)."""
+        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."""
         self.set(name, value)
 
     def __delitem__(self, name):
+        """Deletes a cookie given a name. Wraps cookielib.CookieJar's remove_cookie_by_name()."""
         remove_cookie_by_name(self, name)
 
     def _find(self, name, domain=None, path=None):
+        """Requests uses this method internally to get cookie values. Takes as args name 
+        and optional domain and path. Returns a cookie.value. If there are conflicting cookies,
+        _find arbitrarily chooses one. See _find_no_duplicates if you want an exception thrown
+        if there are conflicting cookies."""
         for cookie in iter(self):
             if cookie.name == name:
                 if domain is None or cookie.domain == domain:
@@ -174,19 +256,39 @@ class RequestsCookieJar(cookielib.CookieJar, collections.MutableMapping):
 
         raise KeyError('name=%r, domain=%r, path=%r' % (name, domain, path))
 
+    def _find_no_duplicates(self, name, domain=None, path=None):
+        """__get_item__ and get call _find_no_duplicates -- never used in Requests internally. 
+        Takes as args name and optional domain and path. Returns a cookie.value. 
+        Throws KeyError if cookie is not found and CookieConflictError if there are 
+        multiple cookies that match name and optionally domain and path."""
+        toReturn = None
+        for cookie in iter(self):
+            if cookie.name == name:
+                if domain is None or cookie.domain == domain:
+                    if path is None or cookie.path == path:
+                        if toReturn != None: # if there are multiple cookies that meet passed in criteria
+                            raise CookieConflictError('There are multiple cookies with name, %r' % (name))
+                        toReturn = cookie.value # we will eventually return this as long as no cookie conflict
+
+        if toReturn:
+            return toReturn
+        raise KeyError('name=%r, domain=%r, path=%r' % (name, domain, path))
+
     def __getstate__(self):
+        """Unlike a normal CookieJar, this class is pickleable."""
         state = self.__dict__.copy()
         # remove the unpickleable RLock object
         state.pop('_cookies_lock')
         return state
 
     def __setstate__(self, state):
+        """Unlike a normal CookieJar, this class is pickleable."""
         self.__dict__.update(state)
         if '_cookies_lock' not in self.__dict__:
             self._cookies_lock = threading.RLock()
 
     def copy(self):
-        """We're probably better off forbidding this."""
+        """This is not implemented. Calling this will throw an exception."""
         raise NotImplementedError
 
 def create_cookie(name, value, **kwargs):

tests/test_cookies.py

@@ -24,6 +24,7 @@ class CookieTests(TestBaseMixin, unittest.TestCase):
 
         # test deprecated dictionary interface
         self.assertEqual(r.cookies['myname'], 'myvalue')
+        self.assertEqual(r.cookies.get('myname'), 'myvalue')
         # test CookieJar interface
         jar = r.cookies
         self.assertEqual(len(jar), 1)
@@ -122,6 +123,38 @@ class CookieTests(TestBaseMixin, unittest.TestCase):
         r = s.get(httpbin('cookies'))
         self.assertEqual(json.loads(r.text)['cookies'], {})
 
+    def test_jar_utility_functions(self):
+        """Test utility functions such as list_domains, list_paths, multiple_domains."""
+        r = requests.get("http://github.com")
+        c = r.cookies
+        # github should send us cookies
+        self.assertTrue(len(c) >= 1)
+        self.assertEqual(len(c), len(r.cookies.keys()))
+        self.assertEqual(len(c), len(r.cookies.values()))
+        self.assertEqual(len(c), len(r.cookies.items()))
+        
+        # domain and path utility functions
+        domain = r.cookies.list_domains()[0]
+        path = r.cookies.list_paths()[0]
+        self.assertEqual(dict(r.cookies), r.cookies.get_dict(domain=domain, path=path))
+        self.assertEqual(len(r.cookies.list_domains()), 1)
+        self.assertEqual(len(r.cookies.list_paths()), 1)
+        self.assertFalse(r.cookies.multiple_domains())
+
+    def test_convert_jar_to_dict(self):
+        """Test that keys, values, and items are defined and that we can convert
+        cookie jars to plain old Python dicts."""
+        r = requests.get(httpbin('cookies', 'set', 'myname', 'myvalue'))
+
+        # test keys, values, and items
+        self.assertEqual(r.cookies.keys(), ['myname'])
+        self.assertEqual(r.cookies.values(), ['myvalue'])
+        self.assertEqual(r.cookies.items(), [('myname','myvalue')])
+        
+        # test if we can convert jar to dict
+        dictOfCookies = dict(r.cookies)
+        self.assertEqual(dictOfCookies, {'myname':'myvalue'})
+        self.assertEqual(dictOfCookies, r.cookies.get_dict())
 
 class LWPCookieJarTest(TestBaseMixin, unittest.TestCase):
     """Check store/load of cookies to FileCookieJar's, specifically LWPCookieJar's."""