Merge pull request #385 from twisted/384-mypy-cookiejar

Add `treq.cookies` helper module

A bunch of cookie-related stuff, incrementally stepping towards addressing #325:

- Add tests that cover the (suprising!) way that cookiejars get updated.
- Extract a scoped_cookie() helper for creating Cookie objects.
- Add a search() function for searching a cookiejar.
- Fix a CookieAgent MyPy failure.
- Improve the typing of treq.response._Response (including its cookie() method).

Fixes #384.

changelog.d/384.feature.rst

@@ -0,0 +1 @@
+The new :mod:`treq.cookies` module provides helper functions for working with `http.cookiejar.Cookie` and `CookieJar` objects.

docs/api.rst

@@ -89,6 +89,15 @@ Authentication
 
 .. autoexception:: UnknownAuthConfig
 
+Cookies
+-------
+
+.. module:: treq.cookies
+
+.. autofunction:: scoped_cookie
+
+.. autofunction:: search
+
 Test Helpers
 ------------
 

docs/examples/using_cookies.py

@@ -4,19 +4,16 @@
 import treq
 
 
-def main(reactor, *args):
-    d = treq.get('https://httpbin.org/cookies/set?hello=world')
+async def main(reactor):
+    resp = await treq.get("https://httpbin.org/cookies/set?hello=world")
 
-    def _get_jar(resp):
-        jar = resp.cookies()
+    jar = resp.cookies()
+    [cookie] = treq.cookies.search(jar, domain="httpbin.org", name="hello")
+    print("The server set our hello cookie to: {}".format(cookie.value))
 
-        print('The server set our hello cookie to: {}'.format(jar['hello']))
+    await treq.get("https://httpbin.org/cookies", cookies=jar).addCallback(
+        print_response
+    )
 
-        return treq.get('https://httpbin.org/cookies', cookies=jar)
 
-    d.addCallback(_get_jar)
-    d.addCallback(print_response)
-
-    return d
-
-react(main, [])
+react(main)

pyproject.toml

@@ -14,6 +14,9 @@ directory = "changelog.d"
 title_format = "{version} ({project_date})"
 issue_format = "`#{issue} <https://github.com/twisted/treq/issues/{issue}>`__"
 
+[tool.ruff]
+line-length = 88
+
 [tool.mypy]
 namespace_packages = true
 plugins = "mypy_zope:plugin"

src/treq/client.py

@@ -2,29 +2,53 @@
 import mimetypes
 import uuid
 from collections import abc
-from http.cookiejar import Cookie, CookieJar
+from http.cookiejar import CookieJar
 from json import dumps as json_dumps
-from typing import (Any, Callable, Iterable, Iterator, List, Mapping,
-                    Optional, Tuple, Union)
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+)
 from urllib.parse import quote_plus
 from urllib.parse import urlencode as _urlencode
 
 from hyperlink import DecodedURL, EncodedURL
 from requests.cookies import merge_cookies
+from treq.cookies import scoped_cookie
 from twisted.internet.defer import Deferred
 from twisted.internet.interfaces import IProtocol
 from twisted.python.components import proxyForInterface, registerAdapter
 from twisted.python.filepath import FilePath
-from twisted.web.client import (BrowserLikeRedirectAgent, ContentDecoderAgent,
-                                CookieAgent, FileBodyProducer, GzipDecoder,
-                                IAgent, RedirectAgent)
+from twisted.web.client import (
+    BrowserLikeRedirectAgent,
+    ContentDecoderAgent,
+    CookieAgent,
+    FileBodyProducer,
+    GzipDecoder,
+    IAgent,
+    RedirectAgent,
+)
 from twisted.web.http_headers import Headers
 from twisted.web.iweb import IBodyProducer, IResponse
 
 from treq import multipart
-from treq._types import (_CookiesType, _DataType, _FilesType, _FileValue,
-                         _HeadersType, _ITreqReactor, _JSONType, _ParamsType,
-                         _URLType)
+from treq._types import (
+    _CookiesType,
+    _DataType,
+    _FilesType,
+    _FileValue,
+    _HeadersType,
+    _ITreqReactor,
+    _JSONType,
+    _ParamsType,
+    _URLType,
+)
 from treq.auth import add_auth
 from treq.response import _Response
 
@@ -55,39 +79,7 @@ def _scoped_cookiejar_from_dict(
     if cookie_dict is None:
         return cookie_jar
     for k, v in cookie_dict.items():
-        secure = url_object.scheme == "https"
-        port_specified = not (
-            (url_object.scheme == "https" and url_object.port == 443)
-            or (url_object.scheme == "http" and url_object.port == 80)
-        )
-        port = str(url_object.port) if port_specified else None
-        domain = url_object.host
-        netscape_domain = domain if "." in domain else domain + ".local"
-
-        cookie_jar.set_cookie(
-            Cookie(
-                # Scoping
-                domain=netscape_domain,
-                port=port,
-                secure=secure,
-                port_specified=port_specified,
-                # Contents
-                name=k,
-                value=v,
-                # Constant/always-the-same stuff
-                version=0,
-                path="/",
-                expires=None,
-                discard=False,
-                comment=None,
-                comment_url=None,
-                rfc2109=False,
-                path_specified=False,
-                domain_specified=False,
-                domain_initial_dot=False,
-                rest={},
-            )
-        )
+        cookie_jar.set_cookie(scoped_cookie(url_object, k, v))
     return cookie_jar
 
 
@@ -254,8 +246,8 @@ def request(
         if not isinstance(cookies, CookieJar):
             cookies = _scoped_cookiejar_from_dict(parsed_url, cookies)
 
-        cookies = merge_cookies(self._cookiejar, cookies)
-        wrapped_agent: IAgent = CookieAgent(self._agent, cookies)
+        merge_cookies(self._cookiejar, cookies)
+        wrapped_agent: IAgent = CookieAgent(self._agent, self._cookiejar)
 
         if allow_redirects:
             if browser_like_redirects:
@@ -289,7 +281,7 @@ def gotResult(result):
         if not unbuffered:
             d.addCallback(_BufferedResponse)
 
-        return d.addCallback(_Response, cookies)
+        return d.addCallback(_Response, self._cookiejar)
 
     def _request_headers(
         self, headers: Optional[_HeadersType], stacklevel: int

src/treq/cookies.py

@@ -0,0 +1,99 @@
+"""
+Convenience helpers for :mod:`http.cookiejar`
+"""
+
+from typing import Union, Iterable, Optional
+from http.cookiejar import Cookie, CookieJar
+
+from hyperlink import EncodedURL
+
+
+def scoped_cookie(origin: Union[str, EncodedURL], name: str, value: str) -> Cookie:
+    """
+    Create a cookie scoped to a given URL's origin.
+
+    You can insert the result directly into a `CookieJar`, like::
+
+        jar = CookieJar()
+        jar.set_cookie(scoped_cookie("https://example.tld", "flavor", "chocolate"))
+
+        await treq.get("https://domain.example", cookies=jar)
+
+    :param origin:
+        A URL that specifies the domain and port number of the cookie.
+
+        If the protocol is HTTP*S* the cookie is marked ``Secure``, meaning
+        it will not be attached to HTTP requests. Otherwise the cookie will be
+        attached to both HTTP and HTTPS requests
+
+    :param name: Name of the cookie.
+
+    :param value: Value of the cookie.
+
+    .. note::
+
+        This does not scope the cookies to any particular path, only the
+        host, port, and scheme of the given URL.
+    """
+    if isinstance(origin, EncodedURL):
+        url_object = origin
+    else:
+        url_object = EncodedURL.from_text(origin)
+
+    secure = url_object.scheme == "https"
+    port_specified = not (
+        (url_object.scheme == "https" and url_object.port == 443)
+        or (url_object.scheme == "http" and url_object.port == 80)
+    )
+    port = str(url_object.port) if port_specified else None
+    domain = url_object.host
+    netscape_domain = domain if "." in domain else domain + ".local"
+    return Cookie(
+        # Scoping
+        domain=netscape_domain,
+        port=port,
+        secure=secure,
+        port_specified=port_specified,
+        # Contents
+        name=name,
+        value=value,
+        # Constant/always-the-same stuff
+        version=0,
+        path="/",
+        expires=None,
+        discard=False,
+        comment=None,
+        comment_url=None,
+        rfc2109=False,
+        path_specified=False,
+        domain_specified=False,
+        domain_initial_dot=False,
+        rest={},
+    )
+
+
+def search(
+    jar: CookieJar, *, domain: str, name: Optional[str] = None
+) -> Iterable[Cookie]:
+    """
+    Raid the cookie jar for matching cookies.
+
+    This is O(n) on the number of cookies in the jar.
+
+    :param jar: The `CookieJar` (or subclass thereof) to search.
+
+    :param domain:
+        Domain, as in the URL, to match. ``.local`` is appended to
+        a bare hostname. Subdomains are not matched (i.e., searching
+        for ``foo.bar.tld`` won't return a cookie set for ``bar.tld``).
+
+    :param name: Cookie name to match (exactly)
+    """
+    netscape_domain = domain if "." in domain else domain + ".local"
+
+    for c in jar:
+        if c.domain != netscape_domain:
+            continue
+        if name is not None and c.name != name:
+            continue
+        yield c

src/treq/response.py

@@ -1,4 +1,7 @@
+from typing import Any, Callable, List
 from requests.cookies import cookiejar_from_dict
+from http.cookiejar import CookieJar
+from twisted.internet.defer import Deferred
 from twisted.python import reflect
 from twisted.python.components import proxyForInterface
 from twisted.web.iweb import UNKNOWN_LENGTH, IResponse
@@ -12,11 +15,14 @@ class _Response(proxyForInterface(IResponse)):  # type: ignore
     adds a few convenience methods.
     """
 
-    def __init__(self, original, cookiejar):
+    original: IResponse
+    _cookiejar: CookieJar
+
+    def __init__(self, original: IResponse, cookiejar: CookieJar):
         self.original = original
         self._cookiejar = cookiejar
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """
         Generate a representation of the response which includes the HTTP
         status code, Content-Type header, and body size, if available.
@@ -38,7 +44,7 @@ def __repr__(self):
             size,
         )
 
-    def collect(self, collector):
+    def collect(self, collector: Callable[[bytes], None]) -> "Deferred[None]":
         """
         Incrementally collect the body of the response, per
         :func:`treq.collect()`.
@@ -51,7 +57,7 @@ def collect(self, collector):
         """
         return collect(self.original, collector)
 
-    def content(self):
+    def content(self) -> "Deferred[bytes]":
         """
         Read the entire body all at once, per :func:`treq.content()`.
 
@@ -60,7 +66,7 @@ def content(self):
         """
         return content(self.original)
 
-    def json(self, **kwargs):
+    def json(self, **kwargs: Any) -> "Deferred[Any]":
         """
         Collect the response body as JSON per :func:`treq.json_content()`.
 
@@ -71,7 +77,7 @@ def json(self, **kwargs):
         """
         return json_content(self.original, **kwargs)
 
-    def text(self, encoding="ISO-8859-1"):
+    def text(self, encoding: str = "ISO-8859-1") -> "Deferred[str]":
         """
         Read the entire body all at once as text, per
         :func:`treq.text_content()`.
@@ -81,13 +87,11 @@ def text(self, encoding="ISO-8859-1"):
         """
         return text_content(self.original, encoding)
 
-    def history(self):
+    def history(self) -> "List[_Response]":
         """
         Get a list of all responses that (such as intermediate redirects),
         that ultimately ended in the current response. The responses are
         ordered chronologically.
-
-        :returns: A `list` of :class:`~treq.response._Response` objects
         """
         response = self
         history = []
@@ -99,16 +103,17 @@ def history(self):
         history.reverse()
         return history
 
-    def cookies(self):
+    def cookies(self) -> CookieJar:
         """
         Get a copy of this response's cookies.
-
-        :rtype: :class:`requests.cookies.RequestsCookieJar`
         """
-        jar = cookiejar_from_dict({})
-
-        if self._cookiejar is not None:
-            for cookie in self._cookiejar:
-                jar.set_cookie(cookie)
+        # NB: This actually returns a RequestsCookieJar, but we type it as a
+        # regular CookieJar because we want to ditch requests as a dependency.
+        # Full deprecation deprecation will require a subclass or wrapper that
+        # warns about the RequestCookieJar extensions.
+        jar: CookieJar = cookiejar_from_dict({})
+
+        for cookie in self._cookiejar:
+            jar.set_cookie(cookie)
 
         return jar