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