pferd/PFERD/cookie_jar.py

70 lines
2.1 KiB
Python
Raw Normal View History

2020-04-20 14:38:03 +02:00
"""A helper for requests cookies."""
2020-04-20 13:35:26 +02:00
import logging
from http.cookiejar import LoadError, LWPCookieJar
from pathlib import Path
from typing import Optional
import requests
logger = logging.getLogger(__name__)
2020-04-20 14:38:03 +02:00
2020-04-20 13:35:26 +02:00
class CookieJar:
2020-04-20 14:38:03 +02:00
"""A cookie jar that can be persisted."""
2020-04-20 13:35:26 +02:00
def __init__(self, cookie_file: Optional[Path] = None) -> None:
2020-04-20 14:38:03 +02:00
"""Create a new cookie jar at the given path.
If the path is None, the cookies will not be persisted.
"""
2020-04-20 13:35:26 +02:00
self._cookies: LWPCookieJar
if cookie_file is None:
self._cookies = LWPCookieJar()
else:
self._cookies = LWPCookieJar(cookie_file)
@property
def cookies(self) -> LWPCookieJar:
2020-04-20 14:38:03 +02:00
"""Return the requests cookie jar."""
2020-04-20 13:35:26 +02:00
return self._cookies
def load_cookies(self) -> None:
2020-04-20 14:38:03 +02:00
"""Load all cookies from the file given in the constructor."""
2020-04-20 13:35:26 +02:00
if self._cookies.filename is None:
return
try:
logger.info(f"Loading old cookies from {self._cookies.filename}")
self._cookies.load(ignore_discard=True)
except (FileNotFoundError, LoadError):
logger.warn(
f"No valid cookie file found at {self._cookies.filename}, "
"continuing with no cookies"
)
def save_cookies(self, reason: Optional[str] = None) -> None:
2020-04-20 14:38:03 +02:00
"""Save the cookies in the file given in the constructor."""
2020-04-20 13:35:26 +02:00
if self._cookies.filename is None:
return
if reason is None:
logger.info("Saving cookies")
else:
logger.info(f"Saving cookies ({reason})")
# TODO figure out why ignore_discard is set
# TODO possibly catch a few more exceptions
self._cookies.save(ignore_discard=True)
def create_session(self) -> requests.Session:
2020-04-20 14:38:03 +02:00
"""Create a new session using the cookie jar."""
2020-04-20 13:35:26 +02:00
sess = requests.Session()
# From the request docs: "All requests code should work out of the box
# with externally provided instances of CookieJar, e.g. LWPCookieJar
# and FileCookieJar."
2020-04-20 14:38:03 +02:00
sess.cookies = self.cookies # type: ignore
2020-04-20 13:35:26 +02:00
return sess