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
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 key = old_listener_to_key[listen_ref]
116 dispatch_reg = _key_to_collection[key]
117 if newowner in dispatch_reg:
118 assert dispatch_reg[newowner] == listen_ref
119 else:
120 dispatch_reg[newowner] = listen_ref
121
122 new_listener_to_key[listen_ref] = key
123
124
125def _clear(owner, elements):
126 if not elements:
127 return
128
129 owner = owner.ref
130 listener_to_key = _collection_to_key[owner]
131 for listen_fn in elements:
132 listen_ref = weakref.ref(listen_fn)
133 key = listener_to_key[listen_ref]
134 dispatch_reg = _key_to_collection[key]
135 dispatch_reg.pop(owner, None)
136
137 if not dispatch_reg:
138 del _key_to_collection[key]
139
140
141class _EventKey(object):
142 """Represent :func:`.listen` arguments."""
143
144 __slots__ = (
145 "target",
146 "identifier",
147 "fn",
148 "fn_key",
149 "fn_wrap",
150 "dispatch_target",
151 )
152
153 def __init__(self, target, identifier, fn, dispatch_target, _fn_wrap=None):
154 self.target = target
155 self.identifier = identifier
156 self.fn = fn
157 if isinstance(fn, types.MethodType):
158 self.fn_key = id(fn.__func__), id(fn.__self__)
159 else:
160 self.fn_key = id(fn)
161 self.fn_wrap = _fn_wrap
162 self.dispatch_target = dispatch_target
163
164 @property
165 def _key(self):
166 return (id(self.target), self.identifier, self.fn_key)
167
168 def with_wrapper(self, fn_wrap):
169 if fn_wrap is self._listen_fn:
170 return self
171 else:
172 return _EventKey(
173 self.target,
174 self.identifier,
175 self.fn,
176 self.dispatch_target,
177 _fn_wrap=fn_wrap,
178 )
179
180 def with_dispatch_target(self, dispatch_target):
181 if dispatch_target is self.dispatch_target:
182 return self
183 else:
184 return _EventKey(
185 self.target,
186 self.identifier,
187 self.fn,
188 dispatch_target,
189 _fn_wrap=self.fn_wrap,
190 )
191
192 def listen(self, *args, **kw):
193 once = kw.pop("once", False)
194 once_unless_exception = kw.pop("_once_unless_exception", False)
195 named = kw.pop("named", False)
196
197 target, identifier, fn = (
198 self.dispatch_target,
199 self.identifier,
200 self._listen_fn,
201 )
202
203 dispatch_collection = getattr(target.dispatch, identifier)
204
205 adjusted_fn = dispatch_collection._adjust_fn_spec(fn, named)
206
207 self = self.with_wrapper(adjusted_fn)
208
209 stub_function = getattr(
210 self.dispatch_target.dispatch._events, self.identifier
211 )
212 if hasattr(stub_function, "_sa_warn"):
213 stub_function._sa_warn()
214
215 if once or once_unless_exception:
216 self.with_wrapper(
217 util.only_once(
218 self._listen_fn, retry_on_exception=once_unless_exception
219 )
220 ).listen(*args, **kw)
221 else:
222 self.dispatch_target.dispatch._listen(self, *args, **kw)
223
224 def remove(self):
225 key = self._key
226
227 if key not in _key_to_collection:
228 raise exc.InvalidRequestError(
229 "No listeners found for event %s / %r / %s "
230 % (self.target, self.identifier, self.fn)
231 )
232 dispatch_reg = _key_to_collection.pop(key)
233
234 for collection_ref, listener_ref in dispatch_reg.items():
235 collection = collection_ref()
236 listener_fn = listener_ref()
237 if collection is not None and listener_fn is not None:
238 collection.remove(self.with_wrapper(listener_fn))
239
240 def contains(self):
241 """Return True if this event key is registered to listen."""
242 return self._key in _key_to_collection
243
244 def base_listen(
245 self, propagate=False, insert=False, named=False, retval=None
246 ):
247
248 target, identifier = self.dispatch_target, self.identifier
249
250 dispatch_collection = getattr(target.dispatch, identifier)
251
252 if insert:
253 dispatch_collection.for_modify(target.dispatch).insert(
254 self, propagate
255 )
256 else:
257 dispatch_collection.for_modify(target.dispatch).append(
258 self, propagate
259 )
260
261 @property
262 def _listen_fn(self):
263 return self.fn_wrap or self.fn
264
265 def append_to_list(self, owner, list_):
266 if _stored_in_collection(self, owner):
267 list_.append(self._listen_fn)
268 return True
269 else:
270 return False
271
272 def remove_from_list(self, owner, list_):
273 _removed_from_collection(self, owner)
274 list_.remove(self._listen_fn)
275
276 def prepend_to_list(self, owner, list_):
277 if _stored_in_collection(self, owner):
278 list_.appendleft(self._listen_fn)
279 return True
280 else:
281 return False