Não pode escolher mais do que 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 
 

201 linhas
5.0 KiB

  1. import collections.abc
  2. import threading
  3. import immutables
  4. __all__ = ('ContextVar', 'Context', 'Token', 'copy_context')
  5. _NO_DEFAULT = object()
  6. class ContextMeta(type(collections.abc.Mapping)):
  7. # contextvars.Context is not subclassable.
  8. def __new__(mcls, names, bases, dct):
  9. cls = super().__new__(mcls, names, bases, dct)
  10. if cls.__module__ != 'contextvars' or cls.__name__ != 'Context':
  11. raise TypeError("type 'Context' is not an acceptable base type")
  12. return cls
  13. class Context(collections.abc.Mapping, metaclass=ContextMeta):
  14. def __init__(self):
  15. self._data = immutables.Map()
  16. self._prev_context = None
  17. def run(self, callable, *args, **kwargs):
  18. if self._prev_context is not None:
  19. raise RuntimeError(
  20. 'cannot enter context: {} is already entered'.format(self))
  21. self._prev_context = _get_context()
  22. try:
  23. _set_context(self)
  24. return callable(*args, **kwargs)
  25. finally:
  26. _set_context(self._prev_context)
  27. self._prev_context = None
  28. def copy(self):
  29. new = Context()
  30. new._data = self._data
  31. return new
  32. def __getitem__(self, var):
  33. if not isinstance(var, ContextVar):
  34. raise TypeError(
  35. "a ContextVar key was expected, got {!r}".format(var))
  36. return self._data[var]
  37. def __contains__(self, var):
  38. if not isinstance(var, ContextVar):
  39. raise TypeError(
  40. "a ContextVar key was expected, got {!r}".format(var))
  41. return var in self._data
  42. def __len__(self):
  43. return len(self._data)
  44. def __iter__(self):
  45. return iter(self._data)
  46. class ContextVarMeta(type):
  47. # contextvars.ContextVar is not subclassable.
  48. def __new__(mcls, names, bases, dct):
  49. cls = super().__new__(mcls, names, bases, dct)
  50. if cls.__module__ != 'contextvars' or cls.__name__ != 'ContextVar':
  51. raise TypeError("type 'ContextVar' is not an acceptable base type")
  52. return cls
  53. def __getitem__(cls, name):
  54. return
  55. class ContextVar(metaclass=ContextVarMeta):
  56. def __init__(self, name, *, default=_NO_DEFAULT):
  57. if not isinstance(name, str):
  58. raise TypeError("context variable name must be a str")
  59. self._name = name
  60. self._default = default
  61. @property
  62. def name(self):
  63. return self._name
  64. def get(self, default=_NO_DEFAULT):
  65. ctx = _get_context()
  66. try:
  67. return ctx[self]
  68. except KeyError:
  69. pass
  70. if default is not _NO_DEFAULT:
  71. return default
  72. if self._default is not _NO_DEFAULT:
  73. return self._default
  74. raise LookupError
  75. def set(self, value):
  76. ctx = _get_context()
  77. data = ctx._data
  78. try:
  79. old_value = data[self]
  80. except KeyError:
  81. old_value = Token.MISSING
  82. updated_data = data.set(self, value)
  83. ctx._data = updated_data
  84. return Token(ctx, self, old_value)
  85. def reset(self, token):
  86. if token._used:
  87. raise RuntimeError("Token has already been used once")
  88. if token._var is not self:
  89. raise ValueError(
  90. "Token was created by a different ContextVar")
  91. if token._context is not _get_context():
  92. raise ValueError(
  93. "Token was created in a different Context")
  94. ctx = token._context
  95. if token._old_value is Token.MISSING:
  96. ctx._data = ctx._data.delete(token._var)
  97. else:
  98. ctx._data = ctx._data.set(token._var, token._old_value)
  99. token._used = True
  100. def __repr__(self):
  101. r = '<ContextVar name={!r}'.format(self.name)
  102. if self._default is not _NO_DEFAULT:
  103. r += ' default={!r}'.format(self._default)
  104. return r + ' at {:0x}>'.format(id(self))
  105. class TokenMeta(type):
  106. # contextvars.Token is not subclassable.
  107. def __new__(mcls, names, bases, dct):
  108. cls = super().__new__(mcls, names, bases, dct)
  109. if cls.__module__ != 'contextvars' or cls.__name__ != 'Token':
  110. raise TypeError("type 'Token' is not an acceptable base type")
  111. return cls
  112. class Token(metaclass=TokenMeta):
  113. MISSING = object()
  114. def __init__(self, context, var, old_value):
  115. self._context = context
  116. self._var = var
  117. self._old_value = old_value
  118. self._used = False
  119. @property
  120. def var(self):
  121. return self._var
  122. @property
  123. def old_value(self):
  124. return self._old_value
  125. def __repr__(self):
  126. r = '<Token '
  127. if self._used:
  128. r += ' used'
  129. r += ' var={!r} at {:0x}>'.format(self._var, id(self))
  130. return r
  131. def copy_context():
  132. return _get_context().copy()
  133. def _get_context():
  134. ctx = getattr(_state, 'context', None)
  135. if ctx is None:
  136. ctx = Context()
  137. _state.context = ctx
  138. return ctx
  139. def _set_context(ctx):
  140. _state.context = ctx
  141. _state = threading.local()