391 lines
11 KiB
Python
391 lines
11 KiB
Python
|
# event/registry.py
|
||
|
|
# Copyright (C) 2005-2025 the SQLAlchemy authors and contributors
|
||
|
|
# <see AUTHORS file>
|
||
|
|
#
|
||
|
|
# This module is part of SQLAlchemy and is released under
|
||
|
|
# the MIT License: https://www.opensource.org/licenses/mit-license.php
|
||
|
|
|
||
|
|
"""Provides managed registration services on behalf of :func:`.listen`
|
||
|
|
arguments.
|
||
|
|
|
||
|
|
By "managed registration", we mean that event listening functions and
|
||
|
|
other objects can be added to various collections in such a way that their
|
||
|
|
membership in all those collections can be revoked at once, based on
|
||
|
|
an equivalent :class:`._EventKey`.
|
||
|
|
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import collections
|
||
|
|
import types
|
||
|
|
import typing
|
||
|
|
from typing import Any
|
||
|
|
from typing import Callable
|
||
|
|
from typing import cast
|
||
|
|
from typing import Deque
|
||
|
|
from typing import Dict
|
||
|
|
from typing import Generic
|
||
|
|
from typing import Iterable
|
||
|
|
from typing import Optional
|
||
|
|
from typing import Tuple
|
||
|
|
from typing import TypeVar
|
||
|
|
from typing import Union
|
||
|
|
import weakref
|
||
|
|
|
||
|
|
from .. import exc
|
||
|
|
from .. import util
|
||
|
|
|
||
|
|
if typing.TYPE_CHECKING:
|
||
|
|
from .attr import RefCollection
|
||
|
|
from .base import dispatcher
|
||
|
|
|
||
|
|
_ListenerFnType = Callable[..., Any]
|
||
|
|
_ListenerFnKeyType = Union[int, Tuple[int, int]]
|
||
|
|
_EventKeyTupleType = Tuple[int, str, _ListenerFnKeyType]
|
||
|
|
|
||
|
|
|
||
|
|
_ET = TypeVar("_ET", bound="EventTarget")
|
||
|
|
|
||
|
|
|
||
|
|
class EventTarget:
|
||
|
|
"""represents an event target, that is, something we can listen on
|
||
|
|
either with that target as a class or as an instance.
|
||
|
|
|
||
|
|
Examples include: Connection, Mapper, Table, Session,
|
||
|
|
InstrumentedAttribute, Engine, Pool, Dialect.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
__slots__ = ()
|
||
|
|
|
||
|
|
dispatch: dispatcher[Any]
|
||
|
|
|
||
|
|
|
||
|
|
_RefCollectionToListenerType = Dict[
|
||
|
|
"weakref.ref[RefCollection[Any]]",
|
||
|
|
"weakref.ref[_ListenerFnType]",
|
||
|
|
]
|
||
|
|
|
||
|
|
_key_to_collection: Dict[_EventKeyTupleType, _RefCollectionToListenerType] = (
|
||
|
|
collections.defaultdict(dict)
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
Given an original listen() argument, can locate all
|
||
|
|
listener collections and the listener fn contained
|
||
|
|
|
||
|
|
(target, identifier, fn) -> {
|
||
|
|
ref(listenercollection) -> ref(listener_fn)
|
||
|
|
ref(listenercollection) -> ref(listener_fn)
|
||
|
|
ref(listenercollection) -> ref(listener_fn)
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
|
||
|
|
_ListenerToEventKeyType = Dict[
|
||
|
|
"weakref.ref[_ListenerFnType]",
|
||
|
|
_EventKeyTupleType,
|
||
|
|
]
|
||
|
|
_collection_to_key: Dict[
|
||
|
|
weakref.ref[RefCollection[Any]],
|
||
|
|
_ListenerToEventKeyType,
|
||
|
|
] = collections.defaultdict(dict)
|
||
|
|
"""
|
||
|
|
Given a _ListenerCollection or _ClsLevelListener, can locate
|
||
|
|
all the original listen() arguments and the listener fn contained
|
||
|
|
|
||
|
|
ref(listenercollection) -> {
|
||
|
|
ref(listener_fn) -> (target, identifier, fn),
|
||
|
|
ref(listener_fn) -> (target, identifier, fn),
|
||
|
|
ref(listener_fn) -> (target, identifier, fn),
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
def _collection_gced(ref: weakref.ref[Any]) -> None:
|
||
|
|
# defaultdict, so can't get a KeyError
|
||
|
|
if not _collection_to_key or ref not in _collection_to_key:
|
||
|
|
return
|
||
|
|
|
||
|
|
ref = cast("weakref.ref[RefCollection[EventTarget]]", ref)
|
||
|
|
|
||
|
|
listener_to_key = _collection_to_key.pop(ref)
|
||
|
|
for key in listener_to_key.values():
|
||
|
|
if key in _key_to_collection:
|
||
|
|
# defaultdict, so can't get a KeyError
|
||
|
|
dispatch_reg = _key_to_collection[key]
|
||
|
|
dispatch_reg.pop(ref)
|
||
|
|
if not dispatch_reg:
|
||
|
|
_key_to_collection.pop(key)
|
||
|
|
|
||
|
|
|
||
|
|
def _stored_in_collection(
|
||
|
|
event_key: _EventKey[_ET], owner: RefCollection[_ET]
|
||
|
|
) -> bool:
|
||
|
|
key = event_key._key
|
||
|
|
|
||
|
|
dispatch_reg = _key_to_collection[key]
|
||
|
|
|
||
|
|
owner_ref = owner.ref
|
||
|
|
listen_ref = weakref.ref(event_key._listen_fn)
|
||
|
|
|
||
|
|
if owner_ref in dispatch_reg:
|
||
|
|
return False
|
||
|
|
|
||
|
|
dispatch_reg[owner_ref] = listen_ref
|
||
|
|
|
||
|
|
listener_to_key = _collection_to_key[owner_ref]
|
||
|
|
listener_to_key[listen_ref] = key
|
||
|
|
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def _removed_from_collection(
|
||
|
|
event_key: _EventKey[_ET], owner: RefCollection[_ET]
|
||
|
|
) -> None:
|
||
|
|
key = event_key._key
|
||
|
|
|
||
|
|
dispatch_reg = _key_to_collection[key]
|
||
|
|
|
||
|
|
listen_ref = weakref.ref(event_key._listen_fn)
|
||
|
|
|
||
|
|
owner_ref = owner.ref
|
||
|
|
dispatch_reg.pop(owner_ref, None)
|
||
|
|
if not dispatch_reg:
|
||
|
|
del _key_to_collection[key]
|
||
|
|
|
||
|
|
if owner_ref in _collection_to_key:
|
||
|
|
listener_to_key = _collection_to_key[owner_ref]
|
||
|
|
# see #12216 - this guards against a removal that already occurred
|
||
|
|
# here. however, I cannot come up with a test that shows any negative
|
||
|
|
# side effects occurring from this removal happening, even though an
|
||
|
|
# event key may still be referenced from a clsleveldispatch here
|
||
|
|
listener_to_key.pop(listen_ref, None)
|
||
|
|
|
||
|
|
|
||
|
|
def _stored_in_collection_multi(
|
||
|
|
newowner: RefCollection[_ET],
|
||
|
|
oldowner: RefCollection[_ET],
|
||
|
|
elements: Iterable[_ListenerFnType],
|
||
|
|
) -> None:
|
||
|
|
if not elements:
|
||
|
|
return
|
||
|
|
|
||
|
|
oldowner_ref = oldowner.ref
|
||
|
|
newowner_ref = newowner.ref
|
||
|
|
|
||
|
|
old_listener_to_key = _collection_to_key[oldowner_ref]
|
||
|
|
new_listener_to_key = _collection_to_key[newowner_ref]
|
||
|
|
|
||
|
|
for listen_fn in elements:
|
||
|
|
listen_ref = weakref.ref(listen_fn)
|
||
|
|
try:
|
||
|
|
key = old_listener_to_key[listen_ref]
|
||
|
|
except KeyError:
|
||
|
|
# can occur during interpreter shutdown.
|
||
|
|
# see #6740
|
||
|
|
continue
|
||
|
|
|
||
|
|
try:
|
||
|
|
dispatch_reg = _key_to_collection[key]
|
||
|
|
except KeyError:
|
||
|
|
continue
|
||
|
|
|
||
|
|
if newowner_ref in dispatch_reg:
|
||
|
|
assert dispatch_reg[newowner_ref] == listen_ref
|
||
|
|
else:
|
||
|
|
dispatch_reg[newowner_ref] = listen_ref
|
||
|
|
|
||
|
|
new_listener_to_key[listen_ref] = key
|
||
|
|
|
||
|
|
|
||
|
|
def _clear(
|
||
|
|
owner: RefCollection[_ET],
|
||
|
|
elements: Iterable[_ListenerFnType],
|
||
|
|
) -> None:
|
||
|
|
if not elements:
|
||
|
|
return
|
||
|
|
|
||
|
|
owner_ref = owner.ref
|
||
|
|
listener_to_key = _collection_to_key[owner_ref]
|
||
|
|
for listen_fn in elements:
|
||
|
|
listen_ref = weakref.ref(listen_fn)
|
||
|
|
key = listener_to_key[listen_ref]
|
||
|
|
dispatch_reg = _key_to_collection[key]
|
||
|
|
dispatch_reg.pop(owner_ref, None)
|
||
|
|
|
||
|
|
if not dispatch_reg:
|
||
|
|
del _key_to_collection[key]
|
||
|
|
|
||
|
|
|
||
|
|
class _EventKey(Generic[_ET]):
|
||
|
|
"""Represent :func:`.listen` arguments."""
|
||
|
|
|
||
|
|
__slots__ = (
|
||
|
|
"target",
|
||
|
|
"identifier",
|
||
|
|
"fn",
|
||
|
|
"fn_key",
|
||
|
|
"fn_wrap",
|
||
|
|
"dispatch_target",
|
||
|
|
)
|
||
|
|
|
||
|
|
target: _ET
|
||
|
|
identifier: str
|
||
|
|
fn: _ListenerFnType
|
||
|
|
fn_key: _ListenerFnKeyType
|
||
|
|
dispatch_target: Any
|
||
|
|
_fn_wrap: Optional[_ListenerFnType]
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
target: _ET,
|
||
|
|
identifier: str,
|
||
|
|
fn: _ListenerFnType,
|
||
|
|
dispatch_target: Any,
|
||
|
|
_fn_wrap: Optional[_ListenerFnType] = None,
|
||
|
|
):
|
||
|
|
self.target = target
|
||
|
|
self.identifier = identifier
|
||
|
|
self.fn = fn
|
||
|
|
if isinstance(fn, types.MethodType):
|
||
|
|
self.fn_key = id(fn.__func__), id(fn.__self__)
|
||
|
|
else:
|
||
|
|
self.fn_key = id(fn)
|
||
|
|
self.fn_wrap = _fn_wrap
|
||
|
|
self.dispatch_target = dispatch_target
|
||
|
|
|
||
|
|
@property
|
||
|
|
def _key(self) -> _EventKeyTupleType:
|
||
|
|
return (id(self.target), self.identifier, self.fn_key)
|
||
|
|
|
||
|
|
def with_wrapper(self, fn_wrap: _ListenerFnType) -> _EventKey[_ET]:
|
||
|
|
if fn_wrap is self._listen_fn:
|
||
|
|
return self
|
||
|
|
else:
|
||
|
|
return _EventKey(
|
||
|
|
self.target,
|
||
|
|
self.identifier,
|
||
|
|
self.fn,
|
||
|
|
self.dispatch_target,
|
||
|
|
_fn_wrap=fn_wrap,
|
||
|
|
)
|
||
|
|
|
||
|
|
def with_dispatch_target(self, dispatch_target: Any) -> _EventKey[_ET]:
|
||
|
|
if dispatch_target is self.dispatch_target:
|
||
|
|
return self
|
||
|
|
else:
|
||
|
|
return _EventKey(
|
||
|
|
self.target,
|
||
|
|
self.identifier,
|
||
|
|
self.fn,
|
||
|
|
dispatch_target,
|
||
|
|
_fn_wrap=self.fn_wrap,
|
||
|
|
)
|
||
|
|
|
||
|
|
def listen(self, *args: Any, **kw: Any) -> None:
|
||
|
|
once = kw.pop("once", False)
|
||
|
|
once_unless_exception = kw.pop("_once_unless_exception", False)
|
||
|
|
named = kw.pop("named", False)
|
||
|
|
|
||
|
|
target, identifier, fn = (
|
||
|
|
self.dispatch_target,
|
||
|
|
self.identifier,
|
||
|
|
self._listen_fn,
|
||
|
|
)
|
||
|
|
|
||
|
|
dispatch_collection = getattr(target.dispatch, identifier)
|
||
|
|
|
||
|
|
adjusted_fn = dispatch_collection._adjust_fn_spec(fn, named)
|
||
|
|
|
||
|
|
self = self.with_wrapper(adjusted_fn)
|
||
|
|
|
||
|
|
stub_function = getattr(
|
||
|
|
self.dispatch_target.dispatch._events, self.identifier
|
||
|
|
)
|
||
|
|
if hasattr(stub_function, "_sa_warn"):
|
||
|
|
stub_function._sa_warn()
|
||
|
|
|
||
|
|
if once or once_unless_exception:
|
||
|
|
self.with_wrapper(
|
||
|
|
util.only_once(
|
||
|
|
self._listen_fn, retry_on_exception=once_unless_exception
|
||
|
|
)
|
||
|
|
).listen(*args, **kw)
|
||
|
|
else:
|
||
|
|
self.dispatch_target.dispatch._listen(self, *args, **kw)
|
||
|
|
|
||
|
|
def remove(self) -> None:
|
||
|
|
key = self._key
|
||
|
|
|
||
|
|
if key not in _key_to_collection:
|
||
|
|
raise exc.InvalidRequestError(
|
||
|
|
"No listeners found for event %s / %r / %s "
|
||
|
|
% (self.target, self.identifier, self.fn)
|
||
|
|
)
|
||
|
|
|
||
|
|
dispatch_reg = _key_to_collection.pop(key)
|
||
|
|
|
||
|
|
for collection_ref, listener_ref in dispatch_reg.items():
|
||
|
|
collection = collection_ref()
|
||
|
|
listener_fn = listener_ref()
|
||
|
|
if collection is not None and listener_fn is not None:
|
||
|
|
collection.remove(self.with_wrapper(listener_fn))
|
||
|
|
|
||
|
|
def contains(self) -> bool:
|
||
|
|
"""Return True if this event key is registered to listen."""
|
||
|
|
return self._key in _key_to_collection
|
||
|
|
|
||
|
|
def base_listen(
|
||
|
|
self,
|
||
|
|
propagate: bool = False,
|
||
|
|
insert: bool = False,
|
||
|
|
named: bool = False,
|
||
|
|
retval: Optional[bool] = None,
|
||
|
|
asyncio: bool = False,
|
||
|
|
) -> None:
|
||
|
|
target, identifier = self.dispatch_target, self.identifier
|
||
|
|
|
||
|
|
dispatch_collection = getattr(target.dispatch, identifier)
|
||
|
|
|
||
|
|
for_modify = dispatch_collection.for_modify(target.dispatch)
|
||
|
|
if asyncio:
|
||
|
|
for_modify._set_asyncio()
|
||
|
|
|
||
|
|
if insert:
|
||
|
|
for_modify.insert(self, propagate)
|
||
|
|
else:
|
||
|
|
for_modify.append(self, propagate)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def _listen_fn(self) -> _ListenerFnType:
|
||
|
|
return self.fn_wrap or self.fn
|
||
|
|
|
||
|
|
def append_to_list(
|
||
|
|
self,
|
||
|
|
owner: RefCollection[_ET],
|
||
|
|
list_: Deque[_ListenerFnType],
|
||
|
|
) -> bool:
|
||
|
|
if _stored_in_collection(self, owner):
|
||
|
|
list_.append(self._listen_fn)
|
||
|
|
return True
|
||
|
|
else:
|
||
|
|
return False
|
||
|
|
|
||
|
|
def remove_from_list(
|
||
|
|
self,
|
||
|
|
owner: RefCollection[_ET],
|
||
|
|
list_: Deque[_ListenerFnType],
|
||
|
|
) -> None:
|
||
|
|
_removed_from_collection(self, owner)
|
||
|
|
list_.remove(self._listen_fn)
|
||
|
|
|
||
|
|
def prepend_to_list(
|
||
|
|
self,
|
||
|
|
owner: RefCollection[_ET],
|
||
|
|
list_: Deque[_ListenerFnType],
|
||
|
|
) -> bool:
|
||
|
|
if _stored_in_collection(self, owner):
|
||
|
|
list_.appendleft(self._listen_fn)
|
||
|
|
return True
|
||
|
|
else:
|
||
|
|
return False
|