threadlocal.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. # This file is dual licensed under the terms of the Apache License, Version
  2. # 2.0, and the MIT License. See the LICENSE file in the root of this
  3. # repository for complete details.
  4. """
  5. Primitives to keep context global but thread (and greenlet) local.
  6. """
  7. from __future__ import absolute_import, division, print_function
  8. import contextlib
  9. import uuid
  10. from structlog._config import BoundLoggerLazyProxy
  11. try:
  12. from greenlet import getcurrent
  13. except ImportError:
  14. from threading import local as ThreadLocal
  15. else:
  16. from weakref import WeakKeyDictionary
  17. class ThreadLocal(object):
  18. """
  19. threading.local() replacement for greenlets.
  20. """
  21. def __init__(self):
  22. self.__dict__["_weakdict"] = WeakKeyDictionary()
  23. def __getattr__(self, name):
  24. key = getcurrent()
  25. try:
  26. return self._weakdict[key][name]
  27. except KeyError:
  28. raise AttributeError(name)
  29. def __setattr__(self, name, val):
  30. key = getcurrent()
  31. self._weakdict.setdefault(key, {})[name] = val
  32. def __delattr__(self, name):
  33. key = getcurrent()
  34. try:
  35. del self._weakdict[key][name]
  36. except KeyError:
  37. raise AttributeError(name)
  38. def wrap_dict(dict_class):
  39. """
  40. Wrap a dict-like class and return the resulting class.
  41. The wrapped class and used to keep global in the current thread.
  42. :param type dict_class: Class used for keeping context.
  43. :rtype: `type`
  44. """
  45. Wrapped = type('WrappedDict-' + str(uuid.uuid4()),
  46. (_ThreadLocalDictWrapper,), {})
  47. Wrapped._tl = ThreadLocal()
  48. Wrapped._dict_class = dict_class
  49. return Wrapped
  50. def as_immutable(logger):
  51. """
  52. Extract the context from a thread local logger into an immutable logger.
  53. :param structlog.BoundLogger logger: A logger with *possibly* thread local
  54. state.
  55. :rtype: :class:`~structlog.BoundLogger` with an immutable context.
  56. """
  57. if isinstance(logger, BoundLoggerLazyProxy):
  58. logger = logger.bind()
  59. try:
  60. ctx = logger._context._tl.dict_.__class__(logger._context._dict)
  61. bl = logger.__class__(
  62. logger._logger,
  63. processors=logger._processors,
  64. context={},
  65. )
  66. bl._context = ctx
  67. return bl
  68. except AttributeError:
  69. return logger
  70. @contextlib.contextmanager
  71. def tmp_bind(logger, **tmp_values):
  72. """
  73. Bind *tmp_values* to *logger* & memorize current state. Rewind afterwards.
  74. """
  75. saved = as_immutable(logger)._context
  76. try:
  77. yield logger.bind(**tmp_values)
  78. finally:
  79. logger._context.clear()
  80. logger._context.update(saved)
  81. class _ThreadLocalDictWrapper(object):
  82. """
  83. Wrap a dict-like class and keep the state *global* but *thread-local*.
  84. Attempts to re-initialize only updates the wrapped dictionary.
  85. Useful for short-lived threaded applications like requests in web app.
  86. Use :func:`wrap` to instantiate and use
  87. :func:`structlog._loggers.BoundLogger.new` to clear the context.
  88. """
  89. def __init__(self, *args, **kw):
  90. """
  91. We cheat. A context dict gets never recreated.
  92. """
  93. if args and isinstance(args[0], self.__class__):
  94. # our state is global, no need to look at args[0] if it's of our
  95. # class
  96. self._dict.update(**kw)
  97. else:
  98. self._dict.update(*args, **kw)
  99. @property
  100. def _dict(self):
  101. """
  102. Return or create and return the current context.
  103. """
  104. try:
  105. return self.__class__._tl.dict_
  106. except AttributeError:
  107. self.__class__._tl.dict_ = self.__class__._dict_class()
  108. return self.__class__._tl.dict_
  109. def __repr__(self):
  110. return '<{0}({1!r})>'.format(self.__class__.__name__, self._dict)
  111. def __eq__(self, other):
  112. # Same class == same dictionary
  113. return self.__class__ == other.__class__
  114. def __ne__(self, other):
  115. return not self.__eq__(other)
  116. # Proxy methods necessary for structlog.
  117. # Dunder methods don't trigger __getattr__ so we need to proxy by hand.
  118. def __iter__(self):
  119. return self._dict.__iter__()
  120. def __setitem__(self, key, value):
  121. self._dict[key] = value
  122. def __delitem__(self, key):
  123. self._dict.__delitem__(key)
  124. def __len__(self):
  125. return self._dict.__len__()
  126. def __getattr__(self, name):
  127. method = getattr(self._dict, name)
  128. return method