49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
# 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)
|