| 1 | from contextvars import ContextVar |
| 2 | from typing import Any, TypeVar, cast, Optional, Dict |
| 3 | |
| 4 | T = TypeVar("T") |
| 5 | |
| 6 | # no mutable default — None is safe |
| 7 | _context_data: ContextVar[Optional[Dict[str, Any]]] = ContextVar("_context_data", default=None) |
| 8 | |
| 9 | |
| 10 | def _ensure_context() -> Dict[str, Any]: |
| 11 | """Make sure a context dict exists, and return it.""" |
| 12 | data = _context_data.get() |
| 13 | if data is None: |
| 14 | data = {} |
| 15 | _context_data.set(data) |
| 16 | return data |
| 17 | |
| 18 | |
| 19 | def set_context_data(key: str, value: Any): |
| 20 | """Set context data for the current async/task context.""" |
| 21 | data = _ensure_context() |
| 22 | if data.get(key) == value: |
| 23 | return |
| 24 | data[key] = value |
| 25 | _context_data.set(data) |
| 26 | |
| 27 | |
| 28 | def delete_context_data(key: str): |
| 29 | """Delete a key from the current async/task context.""" |
| 30 | data = _ensure_context() |
| 31 | if key in data: |
| 32 | del data[key] |
| 33 | _context_data.set(data) |
| 34 | |
| 35 | |
| 36 | def get_context_data(key: Optional[str] = None, default: T = None) -> T: |
| 37 | """Get a key from the current context, or the full dict if key is None.""" |
| 38 | data = _ensure_context() |
| 39 | if key is None: |
| 40 | return cast(T, data) |
| 41 | return cast(T, data.get(key, default)) |
| 42 | |
| 43 | |
| 44 | def clear_context_data(): |
| 45 | """Completely clear the context dict.""" |
| 46 | _context_data.set({}) |