pruncalc/cache.py

import datetime
import lzma
import pathlib
import time
import typing

import cbor2
import httpx
import urllib.parse

if typing.TYPE_CHECKING:
	import httpx._types

ONE_DAY = datetime.timedelta(days=1)

CACHE_DIR = pathlib.Path(__file__).parent / 'cache'
client = httpx.Client(transport=httpx.HTTPTransport(http2=True, retries=3), timeout=5) # retry on ConnectError and ConnectTimeout

@typing.overload
def get(url: str, *, json: typing.Literal[True]=True, headers: httpx._types.HeaderTypes|None=None,
		expiry: datetime.timedelta=datetime.timedelta(minutes=10)) -> typing.Any:
	...
@typing.overload
def get(url: str, *, json: typing.Literal[False], headers: httpx._types.HeaderTypes|None=None,
		expiry: datetime.timedelta=datetime.timedelta(minutes=10)) -> str:
	...
def get(url: str, *, json=True, headers=None, expiry=datetime.timedelta(minutes=10)) -> typing.Any:
	parsed = urllib.parse.urlparse(url)
	assert parsed.hostname is not None
	cache_filename = urllib.parse.quote(parsed.path.removeprefix('/'), safe='')
	if parsed.query:
		cache_filename += '?' + urllib.parse.quote(parsed.query, safe='')
	if json:
		cache_filename += '.cbor'
	cache_filename += '.xz'
	cache_path = CACHE_DIR / parsed.hostname / cache_filename

	try:
		if cache_path.stat().st_mtime > time.time() - expiry.total_seconds(): # less than 10 minutes old
			with lzma.open(cache_path, 'rb') as f:
				if json:
					return cbor2.load(f)
				else:
					return f.read().decode('utf-8')
	except FileNotFoundError:
		pass # fall through

	r = get_with_retries(url, headers)
	if r.status_code == 204:
		raise NoContent

	cache_path.parent.mkdir(parents=True, exist_ok=True)
	with lzma.open(cache_path, 'wb') as f:
		if json:
			data = r.json()
			cbor2.dump(data, f)
		else:
			data = r.text
			f.write(data.encode('utf-8'))
	return data

def get_with_retries(url: str, headers: httpx._types.HeaderTypes|None=None) -> httpx.Response:
	for attempt in range(5):
		try:
			return client.get(url, headers=headers).raise_for_status()
		except httpx.ReadTimeout:
			if attempt == 4:
				raise
			else:
				print(url, 'attempt', attempt+1, 'timed out; retrying...')
	raise AssertionError

class NoContent(Exception):
	pass