You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

288 lines
8.1KB

  1. # event/registry.py
  2. # Copyright (C) 2005-2021 the SQLAlchemy authors and contributors
  3. # <see AUTHORS file>
  4. #
  5. # This module is part of SQLAlchemy and is released under
  6. # the MIT License: http://www.opensource.org/licenses/mit-license.php
  7. """Provides managed registration services on behalf of :func:`.listen`
  8. arguments.
  9. By "managed registration", we mean that event listening functions and
  10. other objects can be added to various collections in such a way that their
  11. membership in all those collections can be revoked at once, based on
  12. an equivalent :class:`._EventKey`.
  13. """
  14. from __future__ import absolute_import
  15. import collections
  16. import types
  17. import weakref
  18. from .. import exc
  19. from .. import util
  20. _key_to_collection = collections.defaultdict(dict)
  21. """
  22. Given an original listen() argument, can locate all
  23. listener collections and the listener fn contained
  24. (target, identifier, fn) -> {
  25. ref(listenercollection) -> ref(listener_fn)
  26. ref(listenercollection) -> ref(listener_fn)
  27. ref(listenercollection) -> ref(listener_fn)
  28. }
  29. """
  30. _collection_to_key = collections.defaultdict(dict)
  31. """
  32. Given a _ListenerCollection or _ClsLevelListener, can locate
  33. all the original listen() arguments and the listener fn contained
  34. ref(listenercollection) -> {
  35. ref(listener_fn) -> (target, identifier, fn),
  36. ref(listener_fn) -> (target, identifier, fn),
  37. ref(listener_fn) -> (target, identifier, fn),
  38. }
  39. """
  40. def _collection_gced(ref):
  41. # defaultdict, so can't get a KeyError
  42. if not _collection_to_key or ref not in _collection_to_key:
  43. return
  44. listener_to_key = _collection_to_key.pop(ref)
  45. for key in listener_to_key.values():
  46. if key in _key_to_collection:
  47. # defaultdict, so can't get a KeyError
  48. dispatch_reg = _key_to_collection[key]
  49. dispatch_reg.pop(ref)
  50. if not dispatch_reg:
  51. _key_to_collection.pop(key)
  52. def _stored_in_collection(event_key, owner):
  53. key = event_key._key
  54. dispatch_reg = _key_to_collection[key]
  55. owner_ref = owner.ref
  56. listen_ref = weakref.ref(event_key._listen_fn)
  57. if owner_ref in dispatch_reg:
  58. return False
  59. dispatch_reg[owner_ref] = listen_ref
  60. listener_to_key = _collection_to_key[owner_ref]
  61. listener_to_key[listen_ref] = key
  62. return True
  63. def _removed_from_collection(event_key, owner):
  64. key = event_key._key
  65. dispatch_reg = _key_to_collection[key]
  66. listen_ref = weakref.ref(event_key._listen_fn)
  67. owner_ref = owner.ref
  68. dispatch_reg.pop(owner_ref, None)
  69. if not dispatch_reg:
  70. del _key_to_collection[key]
  71. if owner_ref in _collection_to_key:
  72. listener_to_key = _collection_to_key[owner_ref]
  73. listener_to_key.pop(listen_ref)
  74. def _stored_in_collection_multi(newowner, oldowner, elements):
  75. if not elements:
  76. return
  77. oldowner = oldowner.ref
  78. newowner = newowner.ref
  79. old_listener_to_key = _collection_to_key[oldowner]
  80. new_listener_to_key = _collection_to_key[newowner]
  81. for listen_fn in elements:
  82. listen_ref = weakref.ref(listen_fn)
  83. key = old_listener_to_key[listen_ref]
  84. dispatch_reg = _key_to_collection[key]
  85. if newowner in dispatch_reg:
  86. assert dispatch_reg[newowner] == listen_ref
  87. else:
  88. dispatch_reg[newowner] = listen_ref
  89. new_listener_to_key[listen_ref] = key
  90. def _clear(owner, elements):
  91. if not elements:
  92. return
  93. owner = owner.ref
  94. listener_to_key = _collection_to_key[owner]
  95. for listen_fn in elements:
  96. listen_ref = weakref.ref(listen_fn)
  97. key = listener_to_key[listen_ref]
  98. dispatch_reg = _key_to_collection[key]
  99. dispatch_reg.pop(owner, None)
  100. if not dispatch_reg:
  101. del _key_to_collection[key]
  102. class _EventKey(object):
  103. """Represent :func:`.listen` arguments."""
  104. __slots__ = (
  105. "target",
  106. "identifier",
  107. "fn",
  108. "fn_key",
  109. "fn_wrap",
  110. "dispatch_target",
  111. )
  112. def __init__(self, target, identifier, fn, dispatch_target, _fn_wrap=None):
  113. self.target = target
  114. self.identifier = identifier
  115. self.fn = fn
  116. if isinstance(fn, types.MethodType):
  117. self.fn_key = id(fn.__func__), id(fn.__self__)
  118. else:
  119. self.fn_key = id(fn)
  120. self.fn_wrap = _fn_wrap
  121. self.dispatch_target = dispatch_target
  122. @property
  123. def _key(self):
  124. return (id(self.target), self.identifier, self.fn_key)
  125. def with_wrapper(self, fn_wrap):
  126. if fn_wrap is self._listen_fn:
  127. return self
  128. else:
  129. return _EventKey(
  130. self.target,
  131. self.identifier,
  132. self.fn,
  133. self.dispatch_target,
  134. _fn_wrap=fn_wrap,
  135. )
  136. def with_dispatch_target(self, dispatch_target):
  137. if dispatch_target is self.dispatch_target:
  138. return self
  139. else:
  140. return _EventKey(
  141. self.target,
  142. self.identifier,
  143. self.fn,
  144. dispatch_target,
  145. _fn_wrap=self.fn_wrap,
  146. )
  147. def listen(self, *args, **kw):
  148. once = kw.pop("once", False)
  149. once_unless_exception = kw.pop("_once_unless_exception", False)
  150. named = kw.pop("named", False)
  151. target, identifier, fn = (
  152. self.dispatch_target,
  153. self.identifier,
  154. self._listen_fn,
  155. )
  156. dispatch_collection = getattr(target.dispatch, identifier)
  157. adjusted_fn = dispatch_collection._adjust_fn_spec(fn, named)
  158. self = self.with_wrapper(adjusted_fn)
  159. stub_function = getattr(
  160. self.dispatch_target.dispatch._events, self.identifier
  161. )
  162. if hasattr(stub_function, "_sa_warn"):
  163. stub_function._sa_warn()
  164. if once or once_unless_exception:
  165. self.with_wrapper(
  166. util.only_once(
  167. self._listen_fn, retry_on_exception=once_unless_exception
  168. )
  169. ).listen(*args, **kw)
  170. else:
  171. self.dispatch_target.dispatch._listen(self, *args, **kw)
  172. def remove(self):
  173. key = self._key
  174. if key not in _key_to_collection:
  175. raise exc.InvalidRequestError(
  176. "No listeners found for event %s / %r / %s "
  177. % (self.target, self.identifier, self.fn)
  178. )
  179. dispatch_reg = _key_to_collection.pop(key)
  180. for collection_ref, listener_ref in dispatch_reg.items():
  181. collection = collection_ref()
  182. listener_fn = listener_ref()
  183. if collection is not None and listener_fn is not None:
  184. collection.remove(self.with_wrapper(listener_fn))
  185. def contains(self):
  186. """Return True if this event key is registered to listen."""
  187. return self._key in _key_to_collection
  188. def base_listen(
  189. self,
  190. propagate=False,
  191. insert=False,
  192. named=False,
  193. retval=None,
  194. asyncio=False,
  195. ):
  196. target, identifier = self.dispatch_target, self.identifier
  197. dispatch_collection = getattr(target.dispatch, identifier)
  198. for_modify = dispatch_collection.for_modify(target.dispatch)
  199. if asyncio:
  200. for_modify._set_asyncio()
  201. if insert:
  202. for_modify.insert(self, propagate)
  203. else:
  204. for_modify.append(self, propagate)
  205. @property
  206. def _listen_fn(self):
  207. return self.fn_wrap or self.fn
  208. def append_to_list(self, owner, list_):
  209. if _stored_in_collection(self, owner):
  210. list_.append(self._listen_fn)
  211. return True
  212. else:
  213. return False
  214. def remove_from_list(self, owner, list_):
  215. _removed_from_collection(self, owner)
  216. list_.remove(self._listen_fn)
  217. def prepend_to_list(self, owner, list_):
  218. if _stored_in_collection(self, owner):
  219. list_.appendleft(self._listen_fn)
  220. return True
  221. else:
  222. return False