1# event/registry.py
2# Copyright (C) 2005-2024 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
8"""Provides managed registration services on behalf of :func:`.listen`
9arguments.
10
11By "managed registration", we mean that event listening functions and
12other objects can be added to various collections in such a way that their
13membership in all those collections can be revoked at once, based on
14an equivalent :class:`._EventKey`.
15
16"""
17
18from __future__ import absolute_import
19
20import collections
21import types
22import weakref
23
24from .. import exc
25from .. import util
26
27
28_key_to_collection = collections.defaultdict(dict)
29"""
30Given an original listen() argument, can locate all
31listener collections and the listener fn contained
32
33(target, identifier, fn) -> {
34 ref(listenercollection) -> ref(listener_fn)
35 ref(listenercollection) -> ref(listener_fn)
36 ref(listenercollection) -> ref(listener_fn)
37 }
38"""
39
40_collection_to_key = collections.defaultdict(dict)
41"""
42Given a _ListenerCollection or _ClsLevelListener, can locate
43all the original listen() arguments and the listener fn contained
44
45ref(listenercollection) -> {
46 ref(listener_fn) -> (target, identifier, fn),
47 ref(listener_fn) -> (target, identifier, fn),
48 ref(listener_fn) -> (target, identifier, fn),
49 }
50"""
51
52
53def _collection_gced(ref):
54 # defaultdict, so can't get a KeyError
55 if not _collection_to_key or ref not in _collection_to_key:
56 return
57 listener_to_key = _collection_to_key.pop(ref)
58 for key in listener_to_key.values():
59 if key in _key_to_collection:
60 # defaultdict, so can't get a KeyError
61 dispatch_reg = _key_to_collection[key]
62 dispatch_reg.pop(ref)
63 if not dispatch_reg:
64 _key_to_collection.pop(key)
65
66
67def _stored_in_collection(event_key, owner):
68 key = event_key._key
69
70 dispatch_reg = _key_to_collection[key]
71
72 owner_ref = owner.ref
73 listen_ref = weakref.ref(event_key._listen_fn)
74
75 if owner_ref in dispatch_reg:
76 return False
77
78 dispatch_reg[owner_ref] = listen_ref
79
80 listener_to_key = _collection_to_key[owner_ref]
81 listener_to_key[listen_ref] = key
82
83 return True
84
85
86def _removed_from_collection(event_key, owner):
87 key = event_key._key
88
89 dispatch_reg = _key_to_collection[key]
90
91 listen_ref = weakref.ref(event_key._listen_fn)
92
93 owner_ref = owner.ref
94 dispatch_reg.pop(owner_ref, None)
95 if not dispatch_reg:
96 del _key_to_collection[key]
97
98 if owner_ref in _collection_to_key:
99 listener_to_key = _collection_to_key[owner_ref]
100 listener_to_key.pop(listen_ref)
101
102
103def _stored_in_collection_multi(newowner, oldowner, elements):
104 if not elements:
105 return
106
107 oldowner = oldowner.ref
108 newowner = newowner.ref
109
110 old_listener_to_key = _collection_to_key[oldowner]
111 new_listener_to_key = _collection_to_key[newowner]
112
113 for listen_fn in elements:
114 listen_ref = weakref.ref(listen_fn)
115 try:
116 key = old_listener_to_key[listen_ref]
117 except KeyError:
118 # can occur during interpreter shutdown.
119 # see #6740
120 continue
121
122 try:
123 dispatch_reg = _key_to_collection[key]
124 except KeyError:
125 continue
126
127 if newowner in dispatch_reg:
128 assert dispatch_reg[newowner] == listen_ref
129 else:
130 dispatch_reg[newowner] = listen_ref
131
132 new_listener_to_key[listen_ref] = key
133
134
135def _clear(owner, elements):
136 if not elements:
137 return
138
139 owner = owner.ref
140 listener_to_key = _collection_to_key[owner]
141 for listen_fn in elements:
142 listen_ref = weakref.ref(listen_fn)
143 key = listener_to_key[listen_ref]
144 dispatch_reg = _key_to_collection[key]
145 dispatch_reg.pop(owner, None)
146
147 if not dispatch_reg:
148 del _key_to_collection[key]
149
150
151class _EventKey(object):
152 """Represent :func:`.listen` arguments."""
153
154 __slots__ = (
155 "target",
156 "identifier",
157 "fn",
158 "fn_key",
159 "fn_wrap",
160 "dispatch_target",
161 )
162
163 def __init__(self, target, identifier, fn, dispatch_target, _fn_wrap=None):
164 self.target = target
165 self.identifier = identifier
166 self.fn = fn
167 if isinstance(fn, types.MethodType):
168 self.fn_key = id(fn.__func__), id(fn.__self__)
169 else:
170 self.fn_key = id(fn)
171 self.fn_wrap = _fn_wrap
172 self.dispatch_target = dispatch_target
173
174 @property
175 def _key(self):
176 return (id(self.target), self.identifier, self.fn_key)
177
178 def with_wrapper(self, fn_wrap):
179 if fn_wrap is self._listen_fn:
180 return self
181 else:
182 return _EventKey(
183 self.target,
184 self.identifier,
185 self.fn,
186 self.dispatch_target,
187 _fn_wrap=fn_wrap,
188 )
189
190 def with_dispatch_target(self, dispatch_target):
191 if dispatch_target is self.dispatch_target:
192 return self
193 else:
194 return _EventKey(
195 self.target,
196 self.identifier,
197 self.fn,
198 dispatch_target,
199 _fn_wrap=self.fn_wrap,
200 )
201
202 def listen(self, *args, **kw):
203 once = kw.pop("once", False)
204 once_unless_exception = kw.pop("_once_unless_exception", False)
205 named = kw.pop("named", False)
206
207 target, identifier, fn = (
208 self.dispatch_target,
209 self.identifier,
210 self._listen_fn,
211 )
212
213 dispatch_collection = getattr(target.dispatch, identifier)
214
215 adjusted_fn = dispatch_collection._adjust_fn_spec(fn, named)
216
217 self = self.with_wrapper(adjusted_fn)
218
219 stub_function = getattr(
220 self.dispatch_target.dispatch._events, self.identifier
221 )
222 if hasattr(stub_function, "_sa_warn"):
223 stub_function._sa_warn()
224
225 if once or once_unless_exception:
226 self.with_wrapper(
227 util.only_once(
228 self._listen_fn, retry_on_exception=once_unless_exception
229 )
230 ).listen(*args, **kw)
231 else:
232 self.dispatch_target.dispatch._listen(self, *args, **kw)
233
234 def remove(self):
235 key = self._key
236
237 if key not in _key_to_collection:
238 raise exc.InvalidRequestError(
239 "No listeners found for event %s / %r / %s "
240 % (self.target, self.identifier, self.fn)
241 )
242
243 dispatch_reg = _key_to_collection.pop(key)
244
245 for collection_ref, listener_ref in dispatch_reg.items():
246 collection = collection_ref()
247 listener_fn = listener_ref()
248 if collection is not None and listener_fn is not None:
249 collection.remove(self.with_wrapper(listener_fn))
250
251 def contains(self):
252 """Return True if this event key is registered to listen."""
253 return self._key in _key_to_collection
254
255 def base_listen(
256 self,
257 propagate=False,
258 insert=False,
259 named=False,
260 retval=None,
261 asyncio=False,
262 ):
263
264 target, identifier = self.dispatch_target, self.identifier
265
266 dispatch_collection = getattr(target.dispatch, identifier)
267
268 for_modify = dispatch_collection.for_modify(target.dispatch)
269 if asyncio:
270 for_modify._set_asyncio()
271
272 if insert:
273 for_modify.insert(self, propagate)
274 else:
275 for_modify.append(self, propagate)
276
277 @property
278 def _listen_fn(self):
279 return self.fn_wrap or self.fn
280
281 def append_to_list(self, owner, list_):
282 if _stored_in_collection(self, owner):
283 list_.append(self._listen_fn)
284 return True
285 else:
286 return False
287
288 def remove_from_list(self, owner, list_):
289 _removed_from_collection(self, owner)
290 list_.remove(self._listen_fn)
291
292 def prepend_to_list(self, owner, list_):
293 if _stored_in_collection(self, owner):
294 list_.appendleft(self._listen_fn)
295 return True
296 else:
297 return False