registry.py 8.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297
  1. # event/registry.py
  2. # Copyright (C) 2005-2022 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: https://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. try:
  84. key = old_listener_to_key[listen_ref]
  85. except KeyError:
  86. # can occur during interpreter shutdown.
  87. # see #6740
  88. continue
  89. try:
  90. dispatch_reg = _key_to_collection[key]
  91. except KeyError:
  92. continue
  93. if newowner in dispatch_reg:
  94. assert dispatch_reg[newowner] == listen_ref
  95. else:
  96. dispatch_reg[newowner] = listen_ref
  97. new_listener_to_key[listen_ref] = key
  98. def _clear(owner, elements):
  99. if not elements:
  100. return
  101. owner = owner.ref
  102. listener_to_key = _collection_to_key[owner]
  103. for listen_fn in elements:
  104. listen_ref = weakref.ref(listen_fn)
  105. key = listener_to_key[listen_ref]
  106. dispatch_reg = _key_to_collection[key]
  107. dispatch_reg.pop(owner, None)
  108. if not dispatch_reg:
  109. del _key_to_collection[key]
  110. class _EventKey(object):
  111. """Represent :func:`.listen` arguments."""
  112. __slots__ = (
  113. "target",
  114. "identifier",
  115. "fn",
  116. "fn_key",
  117. "fn_wrap",
  118. "dispatch_target",
  119. )
  120. def __init__(self, target, identifier, fn, dispatch_target, _fn_wrap=None):
  121. self.target = target
  122. self.identifier = identifier
  123. self.fn = fn
  124. if isinstance(fn, types.MethodType):
  125. self.fn_key = id(fn.__func__), id(fn.__self__)
  126. else:
  127. self.fn_key = id(fn)
  128. self.fn_wrap = _fn_wrap
  129. self.dispatch_target = dispatch_target
  130. @property
  131. def _key(self):
  132. return (id(self.target), self.identifier, self.fn_key)
  133. def with_wrapper(self, fn_wrap):
  134. if fn_wrap is self._listen_fn:
  135. return self
  136. else:
  137. return _EventKey(
  138. self.target,
  139. self.identifier,
  140. self.fn,
  141. self.dispatch_target,
  142. _fn_wrap=fn_wrap,
  143. )
  144. def with_dispatch_target(self, dispatch_target):
  145. if dispatch_target is self.dispatch_target:
  146. return self
  147. else:
  148. return _EventKey(
  149. self.target,
  150. self.identifier,
  151. self.fn,
  152. dispatch_target,
  153. _fn_wrap=self.fn_wrap,
  154. )
  155. def listen(self, *args, **kw):
  156. once = kw.pop("once", False)
  157. once_unless_exception = kw.pop("_once_unless_exception", False)
  158. named = kw.pop("named", False)
  159. target, identifier, fn = (
  160. self.dispatch_target,
  161. self.identifier,
  162. self._listen_fn,
  163. )
  164. dispatch_collection = getattr(target.dispatch, identifier)
  165. adjusted_fn = dispatch_collection._adjust_fn_spec(fn, named)
  166. self = self.with_wrapper(adjusted_fn)
  167. stub_function = getattr(
  168. self.dispatch_target.dispatch._events, self.identifier
  169. )
  170. if hasattr(stub_function, "_sa_warn"):
  171. stub_function._sa_warn()
  172. if once or once_unless_exception:
  173. self.with_wrapper(
  174. util.only_once(
  175. self._listen_fn, retry_on_exception=once_unless_exception
  176. )
  177. ).listen(*args, **kw)
  178. else:
  179. self.dispatch_target.dispatch._listen(self, *args, **kw)
  180. def remove(self):
  181. key = self._key
  182. if key not in _key_to_collection:
  183. raise exc.InvalidRequestError(
  184. "No listeners found for event %s / %r / %s "
  185. % (self.target, self.identifier, self.fn)
  186. )
  187. dispatch_reg = _key_to_collection.pop(key)
  188. for collection_ref, listener_ref in dispatch_reg.items():
  189. collection = collection_ref()
  190. listener_fn = listener_ref()
  191. if collection is not None and listener_fn is not None:
  192. collection.remove(self.with_wrapper(listener_fn))
  193. def contains(self):
  194. """Return True if this event key is registered to listen."""
  195. return self._key in _key_to_collection
  196. def base_listen(
  197. self,
  198. propagate=False,
  199. insert=False,
  200. named=False,
  201. retval=None,
  202. asyncio=False,
  203. ):
  204. target, identifier = self.dispatch_target, self.identifier
  205. dispatch_collection = getattr(target.dispatch, identifier)
  206. for_modify = dispatch_collection.for_modify(target.dispatch)
  207. if asyncio:
  208. for_modify._set_asyncio()
  209. if insert:
  210. for_modify.insert(self, propagate)
  211. else:
  212. for_modify.append(self, propagate)
  213. @property
  214. def _listen_fn(self):
  215. return self.fn_wrap or self.fn
  216. def append_to_list(self, owner, list_):
  217. if _stored_in_collection(self, owner):
  218. list_.append(self._listen_fn)
  219. return True
  220. else:
  221. return False
  222. def remove_from_list(self, owner, list_):
  223. _removed_from_collection(self, owner)
  224. list_.remove(self._listen_fn)
  225. def prepend_to_list(self, owner, list_):
  226. if _stored_in_collection(self, owner):
  227. list_.appendleft(self._listen_fn)
  228. return True
  229. else:
  230. return False