Coverage for stlog/context.py: 100%

1from __future__ import annotations 

2 

3import copy 

4from contextlib import contextmanager 

5from contextvars import ContextVar, Token 

6from typing import Any, Mapping 

7 

8from stlog.base import ( 

9 RESERVED_ATTRS, 

10 StlogError, 

11 check_json_types_or_raise, 

12 get_env_context, 

13) 

14 

15ENV_CONTEXT = get_env_context() 

16_LOGGING_CONTEXT_VAR: ContextVar = ContextVar( 

17 "stlog_logging_context", default=ENV_CONTEXT 

18) 

19 

20 

21class LogContext: 

22 """This is a static class which hosts some utility (static) methods around a "log context" 

23 (global but by execution (worker/thread/async) thanks to contextvars). 

24 

25 All values are (deeply) copied to avoid any changes after adding them to the context. 

26 """ 

27 

28 def __new__(cls): 

29 raise TypeError("This is a static class: do not instanciate it") 

30 

31 @classmethod 

32 def reset_context(cls) -> None: 

33 """Reset the execution log context.""" 

34 _LOGGING_CONTEXT_VAR.set(ENV_CONTEXT) 

35 

36 @classmethod 

37 def _add(cls, **kwargs: Any) -> Token: 

38 for key in kwargs.keys(): 

39 if key in RESERVED_ATTRS: 

40 raise StlogError("key: %s is not allowed (reserved key)" % key) 

41 for val in kwargs.values(): 

42 check_json_types_or_raise(val) 

43 new_context = _LOGGING_CONTEXT_VAR.get() 

44 # we create a new dict here as set() does a shallow copy 

45 return _LOGGING_CONTEXT_VAR.set(copy.deepcopy({**new_context, **kwargs})) 

46 

47 @classmethod 

48 def add(cls, **kwargs: Any) -> None: 

49 """Add some key / values to the execution context. 

50 

51 Only dict, list, int, str, float, bool and None types are allowed 

52 (or composition of these types). 

53 """ 

54 cls._add(**kwargs) 

55 

56 @classmethod 

57 def remove(cls, *keys: str) -> None: 

58 """Remove given keys from the context.""" 

59 # we create a new dict here as set() does a shallow copy 

60 _LOGGING_CONTEXT_VAR.set( 

61 {k: v for k, v in _LOGGING_CONTEXT_VAR.get().items() if k not in keys} 

62 ) 

63 

64 @classmethod 

65 def _get(cls) -> Mapping[str, Any]: 

66 """Get the whole context as a dict.""" 

67 return copy.deepcopy(_LOGGING_CONTEXT_VAR.get()) 

68 

69 @classmethod 

70 def get(cls, key: str, default=None) -> Any: 

71 """Get a context key.""" 

72 return copy.deepcopy(_LOGGING_CONTEXT_VAR.get().get(key, default)) 

73 

74 @classmethod 

75 def getall(cls) -> dict: 

76 """Get the full context as dict.""" 

77 return copy.deepcopy(_LOGGING_CONTEXT_VAR.get()) 

78 

79 @classmethod 

80 @contextmanager 

81 def bind(cls, **kwargs: Any): 

82 """Temporary bind some key / values to the execution log context (context manager).""" 

83 try: 

84 token = cls._add(**kwargs) 

85 yield 

86 finally: 

87 _LOGGING_CONTEXT_VAR.reset(token)