replacer/utils.py

# SPDX-License-Identifier: AGPL-3.0-only

import re
import random
import asyncio
import contextvars
from aiohttp import web
from typing import Dict

class ContextVar:
	"""
	contextvars.ContextVar, but you can call the var to get its value, rather than calling var.get()
	"""
	__slots__ = frozenset({'_var'})
	def __init__(self, *a, **kw): self._var = contextvars.ContextVar(*a, **kw)
	def __call__(self, *args): return self._var.get(*args)  # convenience
	def get(self, *args): return self._var.get(*args)
	def set(self, *args): return self._var.set(*args)
	def reset(self, *args): return self._var.reset(*args)

class MultiReplacer:
	__slots__ = frozenset({'replacements', '_pattern'})

	def __init__(self, replacements: Dict[bytes, bytes]):  # TODO support AnyStr
		self.replacements = replacements
		self._pattern = re.compile(b'(?:%b)' % b'|'.join(map(re.escape, self.replacements)))

	def replace(self, s):
		return self._pattern.sub(self._replacement, s)

	def _replacement(self, match):
		return self.replacements[match[0]]

def build_dprint_factory(debug: bool):
	if debug:
		def dprint_factory():
			# 8 byte pseudo-random request tag
			req_id = ''.join(random.choices('0123456789abcdef', k=8 * 2))
			def dprint(*args, **kwargs): print(f'[{req_id}]', *args, **kwargs)
			return dprint
	else:
		def dprint(*args, **kwargs): return None
		def dprint_factory(): return dprint

	return dprint_factory

def make_unlimited_request(*args):
	return web.BaseRequest(*args, loop=asyncio.get_running_loop(), client_max_size=None)