49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
from contextvars import ContextVar, copy_context
|
|
from typing import Any, Dict, Literal, Union
|
|
|
|
|
|
GUARD_NAME_KEY: Literal["gr.reserved.guard.name"] = "gr.reserved.guard.name"
|
|
DOCUMENT_STORE_KEY: Literal["gr.reserved.document_store"] = "gr.reserved.document_store"
|
|
CALL_KWARGS_KEY: Literal["gr.reserved.call_kwargs"] = "gr.reserved.call_kwargs"
|
|
|
|
|
|
def set_guard_name(guard_name: str) -> None:
|
|
set_context_var(GUARD_NAME_KEY, guard_name)
|
|
|
|
|
|
def get_guard_name() -> str:
|
|
return get_context_var(GUARD_NAME_KEY) or ""
|
|
|
|
|
|
def set_call_kwargs(kwargs: Dict[str, Any]) -> None:
|
|
set_context_var(CALL_KWARGS_KEY, kwargs)
|
|
|
|
|
|
def get_call_kwargs() -> Dict[str, Any]:
|
|
return get_context_var(CALL_KWARGS_KEY) or {}
|
|
|
|
|
|
def get_call_kwarg(kwarg_key: str) -> Union[Any, None]:
|
|
kwargs = get_call_kwargs()
|
|
return kwargs.get(kwarg_key)
|
|
|
|
|
|
def _get_contextvar(key):
|
|
context = copy_context()
|
|
context_var = None
|
|
for c_key in context.keys():
|
|
if c_key.name == key:
|
|
context_var = c_key
|
|
break
|
|
return context_var
|
|
|
|
|
|
def set_context_var(key, value):
|
|
context_var = _get_contextvar(key) or ContextVar(key)
|
|
context_var.set(value)
|
|
|
|
|
|
def get_context_var(key):
|
|
context_var = _get_contextvar(key)
|
|
return context_var.get() if context_var else None
|