pruncalc/cache.py

import lzma
import pathlib
import time
import typing

import cbor2
import httpx
import urllib.parse

if typing.TYPE_CHECKING:
	import httpx._types

CACHE_DIR = pathlib.Path(__file__).parent / 'cache'
client = httpx.Client(transport=httpx.HTTPTransport(http2=True, retries=2), timeout=10)

@typing.overload
def get(url: str, *, json: typing.Literal[True]=True, headers: httpx._types.HeaderTypes|None=None) -> typing.Any:
	...
@typing.overload
def get(url: str, *, json: typing.Literal[False], headers: httpx._types.HeaderTypes|None=None) -> str:
	...
def get(url: str, *, json=True, headers=None) -> typing.Any:
	parsed = urllib.parse.urlparse(url)
	assert parsed.hostname is not None
	cache_filename = urllib.parse.quote(parsed.path.removeprefix('/'), 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() - 600: # 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 = client.get(url, headers=headers).raise_for_status()
	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