From: Tom Christie Date: Fri, 17 May 2019 11:51:00 +0000 (+0100) Subject: Cookie support (#73) X-Git-Tag: 0.3.1~20 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=feae0a895deea88726388a49e22d9e7536330edc;p=thirdparty%2Fhttpx.git Cookie support (#73) * Initial pass at cookie support * Add Cookies model * Support cookie persistence --- diff --git a/httpcore/__init__.py b/httpcore/__init__.py index d30d5365..fddb2cff 100644 --- a/httpcore/__init__.py +++ b/httpcore/__init__.py @@ -5,12 +5,12 @@ from .dispatch.connection import HTTPConnection from .dispatch.connection_pool import ConnectionPool from .exceptions import ( ConnectTimeout, + CookieConflict, DecodingError, InvalidURL, PoolTimeout, ProtocolError, ReadTimeout, - WriteTimeout, RedirectBodyUnavailable, RedirectLoop, ResponseClosed, @@ -18,9 +18,10 @@ from .exceptions import ( StreamConsumed, Timeout, TooManyRedirects, + WriteTimeout, ) from .interfaces import BaseReader, BaseWriter, ConcurrencyBackend, Dispatcher, Protocol -from .models import URL, Headers, Origin, QueryParams, Request, Response +from .models import URL, Cookies, Headers, Origin, QueryParams, Request, Response from .status_codes import codes __version__ = "0.3.0" diff --git a/httpcore/client.py b/httpcore/client.py index 61db5975..3da58e6f 100644 --- a/httpcore/client.py +++ b/httpcore/client.py @@ -18,6 +18,8 @@ from .interfaces import ConcurrencyBackend, Dispatcher from .models import ( URL, AuthTypes, + Cookies, + CookieTypes, Headers, HeaderTypes, QueryParamTypes, @@ -34,6 +36,7 @@ class AsyncClient: def __init__( self, auth: AuthTypes = None, + cookies: CookieTypes = None, ssl: SSLConfig = DEFAULT_SSL_CONFIG, timeout: TimeoutConfig = DEFAULT_TIMEOUT_CONFIG, pool_limits: PoolLimits = DEFAULT_POOL_LIMITS, @@ -47,6 +50,7 @@ class AsyncClient: ) self.auth = auth + self.cookies = Cookies(cookies) self.max_redirects = max_redirects self.dispatch = dispatch @@ -56,6 +60,7 @@ class AsyncClient: *, query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -67,6 +72,7 @@ class AsyncClient: url, query_params=query_params, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -80,6 +86,7 @@ class AsyncClient: *, query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -91,6 +98,7 @@ class AsyncClient: url, query_params=query_params, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -104,6 +112,7 @@ class AsyncClient: *, query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = False, #  Note: Differs to usual default. @@ -115,6 +124,7 @@ class AsyncClient: url, query_params=query_params, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -129,6 +139,7 @@ class AsyncClient: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -141,6 +152,7 @@ class AsyncClient: data=data, query_params=query_params, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -155,6 +167,7 @@ class AsyncClient: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -167,6 +180,7 @@ class AsyncClient: data=data, query_params=query_params, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -181,6 +195,7 @@ class AsyncClient: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -193,6 +208,7 @@ class AsyncClient: data=data, query_params=query_params, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -207,6 +223,7 @@ class AsyncClient: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -219,6 +236,7 @@ class AsyncClient: data=data, query_params=query_params, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -234,6 +252,7 @@ class AsyncClient: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -241,7 +260,12 @@ class AsyncClient: timeout: TimeoutConfig = None, ) -> Response: request = Request( - method, url, data=data, query_params=query_params, headers=headers + method, + url, + data=data, + query_params=query_params, + headers=headers, + cookies=self.merge_cookies(cookies), ) self.prepare_request(request) response = await self.send( @@ -257,6 +281,13 @@ class AsyncClient: def prepare_request(self, request: Request) -> None: request.prepare() + def merge_cookies(self, cookies: CookieTypes = None) -> typing.Optional[CookieTypes]: + if cookies or self.cookies: + merged_cookies = Cookies(self.cookies) + merged_cookies.update(cookies) + return merged_cookies + return cookies + async def send( self, request: Request, @@ -313,6 +344,7 @@ class AsyncClient: request, stream=stream, ssl=ssl, timeout=timeout ) response.history = list(history) + self.cookies.extract_cookies(response) history = [response] + history if not response.is_redirect: break @@ -344,7 +376,8 @@ class AsyncClient: url = self.redirect_url(request, response) headers = self.redirect_headers(request, url) content = self.redirect_content(request, method) - return Request(method=method, url=url, headers=headers, data=content) + cookies = self.merge_cookies(request.cookies) + return Request(method=method, url=url, headers=headers, data=content, cookies=cookies) def redirect_method(self, request: Request, response: Response) -> str: """ @@ -445,6 +478,10 @@ class Client: ) self._loop = asyncio.new_event_loop() + @property + def cookies(self) -> Cookies: + return self._client.cookies + def request( self, method: str, @@ -453,6 +490,7 @@ class Client: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -460,7 +498,12 @@ class Client: timeout: TimeoutConfig = None, ) -> SyncResponse: request = Request( - method, url, data=data, query_params=query_params, headers=headers + method, + url, + data=data, + query_params=query_params, + headers=headers, + cookies=self._client.merge_cookies(cookies), ) self.prepare_request(request) response = self.send( @@ -479,6 +522,7 @@ class Client: *, query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -489,6 +533,7 @@ class Client: "GET", url, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -502,6 +547,7 @@ class Client: *, query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -512,6 +558,7 @@ class Client: "OPTIONS", url, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -525,6 +572,7 @@ class Client: *, query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = False, #  Note: Differs to usual default. @@ -535,6 +583,7 @@ class Client: "HEAD", url, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -549,6 +598,7 @@ class Client: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -560,6 +610,7 @@ class Client: url, data=data, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -574,6 +625,7 @@ class Client: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -585,6 +637,7 @@ class Client: url, data=data, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -599,6 +652,7 @@ class Client: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -610,6 +664,7 @@ class Client: url, data=data, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, @@ -624,6 +679,7 @@ class Client: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, stream: bool = False, auth: AuthTypes = None, allow_redirects: bool = True, @@ -635,6 +691,7 @@ class Client: url, data=data, headers=headers, + cookies=cookies, stream=stream, auth=auth, allow_redirects=allow_redirects, diff --git a/httpcore/exceptions.py b/httpcore/exceptions.py index ad040132..ed311730 100644 --- a/httpcore/exceptions.py +++ b/httpcore/exceptions.py @@ -120,3 +120,9 @@ class InvalidURL(Exception): """ URL was missing a hostname, or was not one of HTTP/HTTPS. """ + + +class CookieConflict(Exception): + """ + Attempted to lookup a cookie by name, but multiple cookies existed. + """ diff --git a/httpcore/models.py b/httpcore/models.py index ea26fa4c..d5b8e7f2 100644 --- a/httpcore/models.py +++ b/httpcore/models.py @@ -1,6 +1,10 @@ import asyncio import cgi +import email.message import typing +import urllib.request +from collections.abc import MutableMapping +from http.cookiejar import Cookie, CookieJar from urllib.parse import parse_qsl, urlencode import chardet @@ -14,6 +18,7 @@ from .decoders import ( MultiDecoder, ) from .exceptions import ( + CookieConflict, HttpError, InvalidURL, ResponseClosed, @@ -43,6 +48,8 @@ HeaderTypes = typing.Union[ typing.List[typing.Tuple[typing.AnyStr, typing.AnyStr]], ] +CookieTypes = typing.Union["Cookies", CookieJar, typing.Dict[str, str]] + AuthTypes = typing.Union[ typing.Tuple[typing.Union[str, bytes], typing.Union[str, bytes]], typing.Callable[["Request"], "Request"], @@ -475,10 +482,14 @@ class Request: data: RequestData = b"", query_params: QueryParamTypes = None, headers: HeaderTypes = None, + cookies: CookieTypes = None, ): self.method = method.upper() self.url = URL(url, query_params=query_params) self.headers = Headers(headers) + if cookies: + self._cookies = Cookies(cookies) + self._cookies.set_cookie_header(self) if isinstance(data, bytes): self.is_streaming = False @@ -536,6 +547,12 @@ class Request: for item in reversed(auto_headers): self.headers.raw.insert(0, item) + @property + def cookies(self) -> "Cookies": + if not hasattr(self, "_cookies"): + self._cookies = Cookies() + return self._cookies + def __repr__(self) -> str: class_name = self.__class__.__name__ url = str(self.url) @@ -756,6 +773,14 @@ class Response: if message: raise HttpError(message) + @property + def cookies(self) -> "Cookies": + if not hasattr(self, "_cookies"): + assert self.request is not None + self._cookies = Cookies() + self._cookies.extract_cookies(self) + return self._cookies + def __repr__(self) -> str: return f"" @@ -835,5 +860,184 @@ class SyncResponse: def close(self) -> None: return self._loop.run_until_complete(self._response.close()) + @property + def cookies(self) -> "Cookies": + return self._response.cookies + def __repr__(self) -> str: return f"" + + +class Cookies(MutableMapping): + """ + HTTP Cookies, as a mutable mapping. + """ + + def __init__(self, cookies: CookieTypes = None) -> None: + if cookies is None or isinstance(cookies, dict): + self.jar = CookieJar() + if isinstance(cookies, dict): + for key, value in cookies.items(): + self.set(key, value) + elif isinstance(cookies, Cookies): + self.jar = CookieJar() + for cookie in cookies.jar: + self.jar.set_cookie(cookie) + else: + self.jar = cookies + + def extract_cookies(self, response: Response) -> None: + """ + Loads any cookies based on the response `Set-Cookie` headers. + """ + assert response.request is not None + urlib_response = self._CookieCompatResponse(response) + urllib_request = self._CookieCompatRequest(response.request) + + self.jar.extract_cookies(urlib_response, urllib_request) # type: ignore + + def set_cookie_header(self, request: Request) -> None: + """ + Sets an appropriate 'Cookie:' HTTP header on the `Request`. + """ + urllib_request = self._CookieCompatRequest(request) + self.jar.add_cookie_header(urllib_request) + + def set(self, name: str, value: str, domain: str = "", path: str = "/") -> None: + """ + Set a cookie value by name. May optionally include domain and path. + """ + kwargs = dict( + version=0, + name=name, + value=value, + port=None, + port_specified=False, + domain=domain, + domain_specified=bool(domain), + domain_initial_dot=domain.startswith("."), + path=path, + path_specified=bool(path), + secure=False, + expires=None, + discard=True, + comment=None, + comment_url=None, + rest={"HttpOnly": None}, + rfc2109=False, + ) + cookie = Cookie(**kwargs) # type: ignore + self.jar.set_cookie(cookie) + + def get( # type: ignore + self, name: str, default: str = None, domain: str = None, path: str = None + ) -> typing.Optional[str]: + """ + Get a cookie by name. May optionally include domain and path + in order to specify exactly which cookie to retrieve. + """ + value = None + for cookie in self.jar: + if cookie.name == name: + if domain is None or cookie.domain == domain: # type: ignore + if path is None or cookie.path == path: + if value is not None: + message = f"Multiple cookies exist with name={name}" + raise CookieConflict(message) + value = cookie.value + + if value is None: + return default + return value + + def delete(self, name: str, domain: str = None, path: str = None) -> None: + """ + Delete a cookie by name. May optionally include domain and path + in order to specify exactly which cookie to delete. + """ + if domain is not None and path is not None: + return self.jar.clear(domain, path, name) + + remove = [] + for cookie in self.jar: + if cookie.name == name: + if domain is None or cookie.domain == domain: # type: ignore + if path is None or cookie.path == path: + remove.append(cookie) + + for cookie in remove: + self.jar.clear(cookie.domain, cookie.path, cookie.name) # type: ignore + + def clear(self, domain: str = None, path: str = None) -> None: + """ + Delete all cookies. Optionally include a domain and path in + order to only delete a subset of all the cookies. + """ + args = [] + if domain is not None: + args.append(domain) + if path is not None: + assert domain is not None + args.append(path) + self.jar.clear(*args) + + def update(self, cookies: CookieTypes) -> None: # type: ignore + cookies = Cookies(cookies) + for cookie in cookies.jar: + self.jar.set_cookie(cookie) + + def __setitem__(self, name: str, value: str) -> None: + return self.set(name, value) + + def __getitem__(self, name: str) -> str: + value = self.get(name) + if value is None: + raise KeyError(name) + return value + + def __delitem__(self, name: str) -> None: + return self.delete(name) + + def __len__(self) -> int: + return len(self.jar) + + def __iter__(self) -> typing.Iterator[str]: + return (cookie.name for cookie in self.jar) + + def __bool__(self) -> bool: + for cookie in self.jar: + return True + return False + + class _CookieCompatRequest(urllib.request.Request): + """ + Wraps a `Request` instance up in a compatability interface suitable + for use with `CookieJar` operations. + """ + + def __init__(self, request: Request) -> None: + super().__init__( + url=str(request.url), + headers=dict(request.headers), + method=request.method, + ) + self.request = request + + def add_unredirected_header(self, key: str, value: str) -> None: + super().add_unredirected_header(key, value) + self.request.headers[key] = value + + class _CookieCompatResponse: + """ + Wraps a `Request` instance up in a compatability interface suitable + for use with `CookieJar` operations. + """ + + def __init__(self, response: Response): + self.response = response + + def info(self) -> email.message.Message: + info = email.message.Message() + for key, value in self.response.headers.items(): + info[key] = value + return info diff --git a/scripts/test b/scripts/test index 221e1d3d..63ad08a6 100755 --- a/scripts/test +++ b/scripts/test @@ -8,5 +8,5 @@ fi set -x -PYTHONPATH=. ${PREFIX}pytest --ignore venv --cov tests --cov ${PACKAGE} --cov-report= --cov-fail-under=100 ${@} -${PREFIX}coverage report -m +PYTHONPATH=. ${PREFIX}pytest --ignore venv --cov tests --cov ${PACKAGE} --cov-report= ${@} +${PREFIX}coverage report --show-missing --fail-under=100 diff --git a/tests/client/test_cookie_handling.py b/tests/client/test_cookie_handling.py new file mode 100644 index 00000000..7fe057d6 --- /dev/null +++ b/tests/client/test_cookie_handling.py @@ -0,0 +1,126 @@ +import json +from http.cookiejar import Cookie, CookieJar + +import pytest + +from httpcore import ( + URL, + Client, + Cookies, + Dispatcher, + Request, + Response, + SSLConfig, + TimeoutConfig, +) + + +class MockDispatch(Dispatcher): + async def send( + self, + request: Request, + stream: bool = False, + ssl: SSLConfig = None, + timeout: TimeoutConfig = None, + ) -> Response: + if request.url.path.startswith("/echo_cookies"): + body = json.dumps({"cookies": request.headers.get("Cookie")}).encode() + return Response(200, content=body, request=request) + elif request.url.path.startswith("/set_cookie"): + headers = {"set-cookie": "example-name=example-value"} + return Response(200, headers=headers, request=request) + + +def test_set_cookie(): + """ + Send a request including a cookie. + """ + url = "http://example.org/echo_cookies" + cookies = {"example-name": "example-value"} + + with Client(dispatch=MockDispatch()) as client: + response = client.get(url, cookies=cookies) + + assert response.status_code == 200 + assert json.loads(response.text) == {"cookies": "example-name=example-value"} + + +def test_set_cookie_with_cookiejar(): + """ + Send a request including a cookie, using a `CookieJar` instance. + """ + + url = "http://example.org/echo_cookies" + cookies = CookieJar() + cookie = Cookie( + version=0, + name="example-name", + value="example-value", + port=None, + port_specified=False, + domain="", + domain_specified=False, + domain_initial_dot=False, + path="/", + path_specified=True, + secure=False, + expires=None, + discard=True, + comment=None, + comment_url=None, + rest={"HttpOnly": None}, + rfc2109=False, + ) + cookies.set_cookie(cookie) + + with Client(dispatch=MockDispatch()) as client: + response = client.get(url, cookies=cookies) + + assert response.status_code == 200 + assert json.loads(response.text) == {"cookies": "example-name=example-value"} + + +def test_set_cookie_with_cookies_model(): + """ + Send a request including a cookie, using a `Cookies` instance. + """ + + url = "http://example.org/echo_cookies" + cookies = Cookies() + cookies["example-name"] = "example-value" + + with Client(dispatch=MockDispatch()) as client: + response = client.get(url, cookies=cookies) + + assert response.status_code == 200 + assert json.loads(response.text) == {"cookies": "example-name=example-value"} + + +def test_get_cookie(): + url = "http://example.org/set_cookie" + + with Client(dispatch=MockDispatch()) as client: + response = client.get(url) + + assert response.status_code == 200 + assert response.cookies["example-name"] == "example-value" + assert client.cookies["example-name"] == "example-value" + + +def test_cookie_persistence(): + """ + Ensure that Client instances persist cookies between requests. + """ + with Client(dispatch=MockDispatch()) as client: + response = client.get("http://example.org/echo_cookies") + assert response.status_code == 200 + assert json.loads(response.text) == {"cookies": None} + + response = client.get("http://example.org/set_cookie") + assert response.status_code == 200 + assert response.cookies["example-name"] == "example-value" + assert client.cookies["example-name"] == "example-value" + + response = client.get("http://example.org/echo_cookies") + assert response.status_code == 200 + assert json.loads(response.text) == {"cookies": "example-name=example-value"} diff --git a/tests/models/test_cookies.py b/tests/models/test_cookies.py new file mode 100644 index 00000000..705245d6 --- /dev/null +++ b/tests/models/test_cookies.py @@ -0,0 +1,50 @@ +import pytest + +from httpcore import CookieConflict, Cookies + + +def test_cookies(): + cookies = Cookies({"name": "value"}) + assert cookies["name"] == "value" + assert "name" in cookies + assert len(cookies) == 1 + assert dict(cookies) == {"name": "value"} + assert bool(cookies) is True + + del cookies["name"] + assert "name" not in cookies + assert len(cookies) == 0 + assert dict(cookies) == {} + assert bool(cookies) is False + + +def test_cookies_update(): + cookies = Cookies() + more_cookies = Cookies() + more_cookies.set("name", "value", domain="example.com") + + cookies.update(more_cookies) + assert dict(cookies) == {"name": "value"} + assert cookies.get("name", domain="example.com") == "value" + + +def test_cookies_with_domain(): + cookies = Cookies() + cookies.set("name", "value", domain="example.com") + cookies.set("name", "value", domain="example.org") + + with pytest.raises(CookieConflict): + cookies["name"] + + cookies.clear(domain="example.com") + assert len(cookies) == 1 + + +def test_cookies_with_domain_and_path(): + cookies = Cookies() + cookies.set("name", "value", domain="example.com", path="/subpath/1") + cookies.set("name", "value", domain="example.com", path="/subpath/2") + cookies.clear(domain="example.com", path="/subpath/1") + assert len(cookies) == 1 + cookies.delete("name", domain="example.com", path="/subpath/2") + assert len(cookies) == 0 diff --git a/tests/test_timeouts.py b/tests/test_timeouts.py index d0cb9b21..f1206fa8 100644 --- a/tests/test_timeouts.py +++ b/tests/test_timeouts.py @@ -6,8 +6,8 @@ from httpcore import ( PoolLimits, PoolTimeout, ReadTimeout, - WriteTimeout, TimeoutConfig, + WriteTimeout, )