Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# orm/collections.py
2# Copyright (C) 2005-2026 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# mypy: allow-untyped-defs, allow-untyped-calls
9"""Support for collections of mapped entities.
11The collections package supplies the machinery used to inform the ORM of
12collection membership changes. An instrumentation via decoration approach is
13used, allowing arbitrary types (including built-ins) to be used as entity
14collections without requiring inheritance from a base class.
16Instrumentation decoration relays membership change events to the
17:class:`.CollectionAttributeImpl` that is currently managing the collection.
18The decorators observe function call arguments and return values, tracking
19entities entering or leaving the collection. Two decorator approaches are
20provided. One is a bundle of generic decorators that map function arguments
21and return values to events::
23 from sqlalchemy.orm.collections import collection
26 class MyClass:
27 # ...
29 @collection.adds(1)
30 def store(self, item):
31 self.data.append(item)
33 @collection.removes_return()
34 def pop(self):
35 return self.data.pop()
37The second approach is a bundle of targeted decorators that wrap appropriate
38append and remove notifiers around the mutation methods present in the
39standard Python ``list``, ``set`` and ``dict`` interfaces. These could be
40specified in terms of generic decorator recipes, but are instead hand-tooled
41for increased efficiency. The targeted decorators occasionally implement
42adapter-like behavior, such as mapping bulk-set methods (``extend``,
43``update``, ``__setslice__``, etc.) into the series of atomic mutation events
44that the ORM requires.
46The targeted decorators are used internally for automatic instrumentation of
47entity collection classes. Every collection class goes through a
48transformation process roughly like so:
501. If the class is a built-in, substitute a trivial sub-class
512. Is this class already instrumented?
523. Add in generic decorators
534. Sniff out the collection interface through duck-typing
545. Add targeted decoration to any undecorated interface method
56This process modifies the class at runtime, decorating methods and adding some
57bookkeeping properties. This isn't possible (or desirable) for built-in
58classes like ``list``, so trivial sub-classes are substituted to hold
59decoration::
61 class InstrumentedList(list):
62 pass
64Collection classes can be specified in ``relationship(collection_class=)`` as
65types or a function that returns an instance. Collection classes are
66inspected and instrumented during the mapper compilation phase. The
67collection_class callable will be executed once to produce a specimen
68instance, and the type of that specimen will be instrumented. Functions that
69return built-in types like ``lists`` will be adapted to produce instrumented
70instances.
72When extending a known type like ``list``, additional decorations are not
73generally not needed. Odds are, the extension method will delegate to a
74method that's already instrumented. For example::
76 class QueueIsh(list):
77 def push(self, item):
78 self.append(item)
80 def shift(self):
81 return self.pop(0)
83There's no need to decorate these methods. ``append`` and ``pop`` are already
84instrumented as part of the ``list`` interface. Decorating them would fire
85duplicate events, which should be avoided.
87The targeted decoration tries not to rely on other methods in the underlying
88collection class, but some are unavoidable. Many depend on 'read' methods
89being present to properly instrument a 'write', for example, ``__setitem__``
90needs ``__getitem__``. "Bulk" methods like ``update`` and ``extend`` may also
91reimplemented in terms of atomic appends and removes, so the ``extend``
92decoration will actually perform many ``append`` operations and not call the
93underlying method at all.
95Tight control over bulk operation and the firing of events is also possible by
96implementing the instrumentation internally in your methods. The basic
97instrumentation package works under the general assumption that collection
98mutation will not raise unusual exceptions. If you want to closely
99orchestrate append and remove events with exception management, internal
100instrumentation may be the answer. Within your method,
101``collection_adapter(self)`` will retrieve an object that you can use for
102explicit control over triggering append and remove events.
104The owning object and :class:`.CollectionAttributeImpl` are also reachable
105through the adapter, allowing for some very sophisticated behavior.
107"""
109from __future__ import annotations
111import operator
112import threading
113import typing
114from typing import Any
115from typing import Callable
116from typing import cast
117from typing import Collection
118from typing import Dict
119from typing import Iterable
120from typing import List
121from typing import NoReturn
122from typing import Optional
123from typing import Protocol
124from typing import Set
125from typing import Tuple
126from typing import Type
127from typing import TYPE_CHECKING
128from typing import TypeVar
129from typing import Union
130import weakref
132from .base import NO_KEY
133from .. import exc as sa_exc
134from .. import util
135from ..sql.base import NO_ARG
136from ..util.compat import inspect_getfullargspec
138if typing.TYPE_CHECKING:
139 from .attributes import _CollectionAttributeImpl
140 from .attributes import AttributeEventToken
141 from .mapped_collection import attribute_keyed_dict
142 from .mapped_collection import column_keyed_dict
143 from .mapped_collection import keyfunc_mapping
144 from .mapped_collection import KeyFuncDict # noqa: F401
145 from .state import InstanceState
148__all__ = [
149 "collection",
150 "collection_adapter",
151 "keyfunc_mapping",
152 "column_keyed_dict",
153 "attribute_keyed_dict",
154 "KeyFuncDict",
155 # old names in < 2.0
156 "mapped_collection",
157 "column_mapped_collection",
158 "attribute_mapped_collection",
159 "MappedCollection",
160]
162__instrumentation_mutex = threading.Lock()
165_CollectionFactoryType = Callable[[], "_AdaptedCollectionProtocol"]
167_T = TypeVar("_T", bound=Any)
168_KT = TypeVar("_KT", bound=Any)
169_VT = TypeVar("_VT", bound=Any)
170_COL = TypeVar("_COL", bound="Collection[Any]")
171_FN = TypeVar("_FN", bound="Callable[..., Any]")
174class _CollectionConverterProtocol(Protocol):
175 def __call__(self, collection: _COL) -> _COL: ...
178class _AdaptedCollectionProtocol(Protocol):
179 _sa_adapter: CollectionAdapter
180 _sa_appender: Callable[..., Any]
181 _sa_remover: Callable[..., Any]
182 _sa_iterator: Callable[..., Iterable[Any]]
185class collection:
186 """Decorators for entity collection classes.
188 The decorators fall into two groups: annotations and interception recipes.
190 The annotating decorators (appender, remover, iterator,
191 internally_instrumented) indicate the method's purpose and take no
192 arguments. They are not written with parens::
194 @collection.appender
195 def append(self, append): ...
197 The recipe decorators all require parens, even those that take no
198 arguments::
200 @collection.adds("entity")
201 def insert(self, position, entity): ...
204 @collection.removes_return()
205 def popitem(self): ...
207 """
209 # Bundled as a class solely for ease of use: packaging, doc strings,
210 # importability.
212 @staticmethod
213 def appender(fn):
214 """Tag the method as the collection appender.
216 The appender method is called with one positional argument: the value
217 to append. The method will be automatically decorated with 'adds(1)'
218 if not already decorated::
220 @collection.appender
221 def add(self, append): ...
224 # or, equivalently
225 @collection.appender
226 @collection.adds(1)
227 def add(self, append): ...
230 # for mapping type, an 'append' may kick out a previous value
231 # that occupies that slot. consider d['a'] = 'foo'- any previous
232 # value in d['a'] is discarded.
233 @collection.appender
234 @collection.replaces(1)
235 def add(self, entity):
236 key = some_key_func(entity)
237 previous = None
238 if key in self:
239 previous = self[key]
240 self[key] = entity
241 return previous
243 If the value to append is not allowed in the collection, you may
244 raise an exception. Something to remember is that the appender
245 will be called for each object mapped by a database query. If the
246 database contains rows that violate your collection semantics, you
247 will need to get creative to fix the problem, as access via the
248 collection will not work.
250 If the appender method is internally instrumented, you must also
251 receive the keyword argument '_sa_initiator' and ensure its
252 promulgation to collection events.
254 """
255 fn._sa_instrument_role = "appender"
256 return fn
258 @staticmethod
259 def remover(fn):
260 """Tag the method as the collection remover.
262 The remover method is called with one positional argument: the value
263 to remove. The method will be automatically decorated with
264 :meth:`removes_return` if not already decorated::
266 @collection.remover
267 def zap(self, entity): ...
270 # or, equivalently
271 @collection.remover
272 @collection.removes_return()
273 def zap(self): ...
275 If the value to remove is not present in the collection, you may
276 raise an exception or return None to ignore the error.
278 If the remove method is internally instrumented, you must also
279 receive the keyword argument '_sa_initiator' and ensure its
280 promulgation to collection events.
282 """
283 fn._sa_instrument_role = "remover"
284 return fn
286 @staticmethod
287 def iterator(fn):
288 """Tag the method as the collection remover.
290 The iterator method is called with no arguments. It is expected to
291 return an iterator over all collection members::
293 @collection.iterator
294 def __iter__(self): ...
296 """
297 fn._sa_instrument_role = "iterator"
298 return fn
300 @staticmethod
301 def internally_instrumented(fn):
302 """Tag the method as instrumented.
304 This tag will prevent any decoration from being applied to the
305 method. Use this if you are orchestrating your own calls to
306 :func:`.collection_adapter` in one of the basic SQLAlchemy
307 interface methods, or to prevent an automatic ABC method
308 decoration from wrapping your implementation::
310 # normally an 'extend' method on a list-like class would be
311 # automatically intercepted and re-implemented in terms of
312 # SQLAlchemy events and append(). your implementation will
313 # never be called, unless:
314 @collection.internally_instrumented
315 def extend(self, items): ...
317 """
318 fn._sa_instrumented = True
319 return fn
321 @staticmethod
322 def adds(arg: int) -> Callable[[_FN], _FN]:
323 """Mark the method as adding an entity to the collection.
325 Adds "add to collection" handling to the method. The decorator
326 argument indicates which method argument holds the SQLAlchemy-relevant
327 value. Arguments can be specified positionally (i.e. integer) or by
328 name::
330 @collection.adds(1)
331 def push(self, item): ...
334 @collection.adds("entity")
335 def do_stuff(self, thing, entity=None): ...
337 """
339 def decorator(fn):
340 fn._sa_instrument_before = ("fire_append_event", arg)
341 return fn
343 return decorator
345 @staticmethod
346 def replaces(arg):
347 """Mark the method as replacing an entity in the collection.
349 Adds "add to collection" and "remove from collection" handling to
350 the method. The decorator argument indicates which method argument
351 holds the SQLAlchemy-relevant value to be added, and return value, if
352 any will be considered the value to remove.
354 Arguments can be specified positionally (i.e. integer) or by name::
356 @collection.replaces(2)
357 def __setitem__(self, index, item): ...
359 """
361 def decorator(fn):
362 fn._sa_instrument_before = ("fire_append_event", arg)
363 fn._sa_instrument_after = "fire_remove_event"
364 return fn
366 return decorator
368 @staticmethod
369 def removes(arg):
370 """Mark the method as removing an entity in the collection.
372 Adds "remove from collection" handling to the method. The decorator
373 argument indicates which method argument holds the SQLAlchemy-relevant
374 value to be removed. Arguments can be specified positionally (i.e.
375 integer) or by name::
377 @collection.removes(1)
378 def zap(self, item): ...
380 For methods where the value to remove is not known at call-time, use
381 collection.removes_return.
383 """
385 def decorator(fn):
386 fn._sa_instrument_before = ("fire_remove_event", arg)
387 return fn
389 return decorator
391 @staticmethod
392 def removes_return():
393 """Mark the method as removing an entity in the collection.
395 Adds "remove from collection" handling to the method. The return
396 value of the method, if any, is considered the value to remove. The
397 method arguments are not inspected::
399 @collection.removes_return()
400 def pop(self): ...
402 For methods where the value to remove is known at call-time, use
403 collection.remove.
405 """
407 def decorator(fn):
408 fn._sa_instrument_after = "fire_remove_event"
409 return fn
411 return decorator
414if TYPE_CHECKING:
416 def collection_adapter(collection: Collection[Any]) -> CollectionAdapter:
417 """Fetch the :class:`.CollectionAdapter` for a collection."""
419else:
420 collection_adapter = operator.attrgetter("_sa_adapter")
423class CollectionAdapter:
424 """Bridges between the ORM and arbitrary Python collections.
426 Proxies base-level collection operations (append, remove, iterate)
427 to the underlying Python collection, and emits add/remove events for
428 entities entering or leaving the collection.
430 The ORM uses :class:`.CollectionAdapter` exclusively for interaction with
431 entity collections.
434 """
436 __slots__ = (
437 "attr",
438 "_key",
439 "_data",
440 "owner_state",
441 "invalidated",
442 "empty",
443 )
445 attr: _CollectionAttributeImpl
446 _key: str
448 # this is actually a weakref; see note in constructor
449 _data: Callable[..., _AdaptedCollectionProtocol]
451 owner_state: InstanceState[Any]
452 invalidated: bool
453 empty: bool
455 def __init__(
456 self,
457 attr: _CollectionAttributeImpl,
458 owner_state: InstanceState[Any],
459 data: _AdaptedCollectionProtocol,
460 ):
461 self.attr = attr
462 self._key = attr.key
464 # this weakref stays referenced throughout the lifespan of
465 # CollectionAdapter. so while the weakref can return None, this
466 # is realistically only during garbage collection of this object, so
467 # we type this as a callable that returns _AdaptedCollectionProtocol
468 # in all cases.
469 self._data = weakref.ref(data) # type: ignore
471 self.owner_state = owner_state
472 data._sa_adapter = self
473 self.invalidated = False
474 self.empty = False
476 def _warn_invalidated(self) -> None:
477 util.warn("This collection has been invalidated.")
479 @property
480 def data(self) -> _AdaptedCollectionProtocol:
481 "The entity collection being adapted."
482 return self._data()
484 @property
485 def _referenced_by_owner(self) -> bool:
486 """return True if the owner state still refers to this collection.
488 This will return False within a bulk replace operation,
489 where this collection is the one being replaced.
491 """
492 return self.owner_state.dict[self._key] is self._data()
494 def bulk_appender(self):
495 return self._data()._sa_appender
497 def append_with_event(
498 self, item: Any, initiator: Optional[AttributeEventToken] = None
499 ) -> None:
500 """Add an entity to the collection, firing mutation events."""
502 self._data()._sa_appender(item, _sa_initiator=initiator)
504 def _set_empty(self, user_data):
505 assert (
506 not self.empty
507 ), "This collection adapter is already in the 'empty' state"
508 self.empty = True
509 self.owner_state._empty_collections[self._key] = user_data
511 def _reset_empty(self) -> None:
512 assert (
513 self.empty
514 ), "This collection adapter is not in the 'empty' state"
515 self.empty = False
516 self.owner_state.dict[self._key] = (
517 self.owner_state._empty_collections.pop(self._key)
518 )
520 def _refuse_empty(self) -> NoReturn:
521 raise sa_exc.InvalidRequestError(
522 "This is a special 'empty' collection which cannot accommodate "
523 "internal mutation operations"
524 )
526 def append_without_event(self, item: Any) -> None:
527 """Add or restore an entity to the collection, firing no events."""
529 if self.empty:
530 self._refuse_empty()
531 self._data()._sa_appender(item, _sa_initiator=False)
533 def append_multiple_without_event(self, items: Iterable[Any]) -> None:
534 """Add or restore an entity to the collection, firing no events."""
535 if self.empty:
536 self._refuse_empty()
537 appender = self._data()._sa_appender
538 for item in items:
539 appender(item, _sa_initiator=False)
541 def bulk_remover(self):
542 return self._data()._sa_remover
544 def remove_with_event(
545 self, item: Any, initiator: Optional[AttributeEventToken] = None
546 ) -> None:
547 """Remove an entity from the collection, firing mutation events."""
548 self._data()._sa_remover(item, _sa_initiator=initiator)
550 def remove_without_event(self, item: Any) -> None:
551 """Remove an entity from the collection, firing no events."""
552 if self.empty:
553 self._refuse_empty()
554 self._data()._sa_remover(item, _sa_initiator=False)
556 def clear_with_event(
557 self, initiator: Optional[AttributeEventToken] = None
558 ) -> None:
559 """Empty the collection, firing a mutation event for each entity."""
561 if self.empty:
562 self._refuse_empty()
563 remover = self._data()._sa_remover
564 for item in list(self):
565 remover(item, _sa_initiator=initiator)
567 def clear_without_event(self) -> None:
568 """Empty the collection, firing no events."""
570 if self.empty:
571 self._refuse_empty()
572 remover = self._data()._sa_remover
573 for item in list(self):
574 remover(item, _sa_initiator=False)
576 def __iter__(self):
577 """Iterate over entities in the collection."""
579 return iter(self._data()._sa_iterator())
581 def __len__(self):
582 """Count entities in the collection."""
583 return len(list(self._data()._sa_iterator()))
585 def __bool__(self):
586 return True
588 def _fire_append_wo_mutation_event_bulk(
589 self, items, initiator=None, key=NO_KEY
590 ):
591 if not items:
592 return
594 if initiator is not False:
595 if self.invalidated:
596 self._warn_invalidated()
598 if self.empty:
599 self._reset_empty()
601 for item in items:
602 self.attr.fire_append_wo_mutation_event(
603 self.owner_state,
604 self.owner_state.dict,
605 item,
606 initiator,
607 key,
608 )
610 def fire_append_wo_mutation_event(self, item, initiator=None, key=NO_KEY):
611 """Notify that a entity is entering the collection but is already
612 present.
615 Initiator is a token owned by the InstrumentedAttribute that
616 initiated the membership mutation, and should be left as None
617 unless you are passing along an initiator value from a chained
618 operation.
620 .. versionadded:: 1.4.15
622 """
623 if initiator is not False:
624 if self.invalidated:
625 self._warn_invalidated()
627 if self.empty:
628 self._reset_empty()
630 return self.attr.fire_append_wo_mutation_event(
631 self.owner_state, self.owner_state.dict, item, initiator, key
632 )
633 else:
634 return item
636 def fire_append_event(self, item, initiator=None, key=NO_KEY):
637 """Notify that a entity has entered the collection.
639 Initiator is a token owned by the InstrumentedAttribute that
640 initiated the membership mutation, and should be left as None
641 unless you are passing along an initiator value from a chained
642 operation.
644 """
645 if initiator is not False:
646 if self.invalidated:
647 self._warn_invalidated()
649 if self.empty:
650 self._reset_empty()
652 return self.attr.fire_append_event(
653 self.owner_state, self.owner_state.dict, item, initiator, key
654 )
655 else:
656 return item
658 def _fire_remove_event_bulk(self, items, initiator=None, key=NO_KEY):
659 if not items:
660 return
662 if initiator is not False:
663 if self.invalidated:
664 self._warn_invalidated()
666 if self.empty:
667 self._reset_empty()
669 for item in items:
670 self.attr.fire_remove_event(
671 self.owner_state,
672 self.owner_state.dict,
673 item,
674 initiator,
675 key,
676 )
678 def fire_remove_event(self, item, initiator=None, key=NO_KEY):
679 """Notify that a entity has been removed from the collection.
681 Initiator is the InstrumentedAttribute that initiated the membership
682 mutation, and should be left as None unless you are passing along
683 an initiator value from a chained operation.
685 """
686 if initiator is not False:
687 if self.invalidated:
688 self._warn_invalidated()
690 if self.empty:
691 self._reset_empty()
693 self.attr.fire_remove_event(
694 self.owner_state, self.owner_state.dict, item, initiator, key
695 )
697 def fire_pre_remove_event(self, initiator=None, key=NO_KEY):
698 """Notify that an entity is about to be removed from the collection.
700 Only called if the entity cannot be removed after calling
701 fire_remove_event().
703 """
704 if self.invalidated:
705 self._warn_invalidated()
706 self.attr.fire_pre_remove_event(
707 self.owner_state,
708 self.owner_state.dict,
709 initiator=initiator,
710 key=key,
711 )
713 def __getstate__(self):
714 return {
715 "key": self._key,
716 "owner_state": self.owner_state,
717 "owner_cls": self.owner_state.class_,
718 "data": self.data,
719 "invalidated": self.invalidated,
720 "empty": self.empty,
721 }
723 def __setstate__(self, d):
724 self._key = d["key"]
725 self.owner_state = d["owner_state"]
727 # see note in constructor regarding this type: ignore
728 self._data = weakref.ref(d["data"]) # type: ignore
730 d["data"]._sa_adapter = self
731 self.invalidated = d["invalidated"]
732 self.attr = getattr(d["owner_cls"], self._key).impl
733 self.empty = d.get("empty", False)
736def bulk_replace(values, existing_adapter, new_adapter, initiator=None):
737 """Load a new collection, firing events based on prior like membership.
739 Appends instances in ``values`` onto the ``new_adapter``. Events will be
740 fired for any instance not present in the ``existing_adapter``. Any
741 instances in ``existing_adapter`` not present in ``values`` will have
742 remove events fired upon them.
744 :param values: An iterable of collection member instances
746 :param existing_adapter: A :class:`.CollectionAdapter` of
747 instances to be replaced
749 :param new_adapter: An empty :class:`.CollectionAdapter`
750 to load with ``values``
753 """
755 assert isinstance(values, list)
757 idset = util.IdentitySet
758 existing_idset = idset(existing_adapter or ())
759 constants = existing_idset.intersection(values or ())
760 additions = idset(values or ()).difference(constants)
761 removals = existing_idset.difference(constants)
763 appender = new_adapter.bulk_appender()
765 for member in values or ():
766 if member in additions:
767 appender(member, _sa_initiator=initiator)
768 elif member in constants:
769 appender(member, _sa_initiator=False)
771 if existing_adapter:
772 existing_adapter._fire_append_wo_mutation_event_bulk(
773 constants, initiator=initiator
774 )
775 existing_adapter._fire_remove_event_bulk(removals, initiator=initiator)
778def _prepare_instrumentation(
779 factory: Union[Type[Collection[Any]], _CollectionFactoryType],
780) -> _CollectionFactoryType:
781 """Prepare a callable for future use as a collection class factory.
783 Given a collection class factory (either a type or no-arg callable),
784 return another factory that will produce compatible instances when
785 called.
787 This function is responsible for converting collection_class=list
788 into the run-time behavior of collection_class=InstrumentedList.
790 """
792 impl_factory: _CollectionFactoryType
794 # Convert a builtin to 'Instrumented*'
795 if factory in __canned_instrumentation:
796 impl_factory = __canned_instrumentation[factory]
797 else:
798 impl_factory = cast(_CollectionFactoryType, factory)
800 cls: Union[_CollectionFactoryType, Type[Collection[Any]]]
802 # Create a specimen
803 cls = type(impl_factory())
805 # Did factory callable return a builtin?
806 if cls in __canned_instrumentation:
807 # if so, just convert.
808 # in previous major releases, this codepath wasn't working and was
809 # not covered by tests. prior to that it supplied a "wrapper"
810 # function that would return the class, though the rationale for this
811 # case is not known
812 impl_factory = __canned_instrumentation[cls]
813 cls = type(impl_factory())
815 # Instrument the class if needed.
816 if __instrumentation_mutex.acquire():
817 try:
818 if getattr(cls, "_sa_instrumented", None) != id(cls):
819 _instrument_class(cls)
820 finally:
821 __instrumentation_mutex.release()
823 return impl_factory
826def _instrument_class(cls):
827 """Modify methods in a class and install instrumentation."""
829 # In the normal call flow, a request for any of the 3 basic collection
830 # types is transformed into one of our trivial subclasses
831 # (e.g. InstrumentedList). Catch anything else that sneaks in here...
832 if cls.__module__ == "__builtin__":
833 raise sa_exc.ArgumentError(
834 "Can not instrument a built-in type. Use a "
835 "subclass, even a trivial one."
836 )
838 roles, methods = _locate_roles_and_methods(cls)
840 _setup_canned_roles(cls, roles, methods)
842 _assert_required_roles(cls, roles, methods)
844 _set_collection_attributes(cls, roles, methods)
847def _locate_roles_and_methods(cls):
848 """search for _sa_instrument_role-decorated methods in
849 method resolution order, assign to roles.
851 """
853 roles: Dict[str, str] = {}
854 methods: Dict[str, Tuple[Optional[str], Optional[int], Optional[str]]] = {}
856 for supercls in cls.__mro__:
857 for name, method in vars(supercls).items():
858 if not callable(method):
859 continue
861 # note role declarations
862 if hasattr(method, "_sa_instrument_role"):
863 role = method._sa_instrument_role
864 assert role in ("appender", "remover", "iterator")
865 roles.setdefault(role, name)
867 # transfer instrumentation requests from decorated function
868 # to the combined queue
869 before: Optional[Tuple[str, int]] = None
870 after: Optional[str] = None
872 if hasattr(method, "_sa_instrument_before"):
873 op, argument = method._sa_instrument_before
874 assert op in ("fire_append_event", "fire_remove_event")
875 before = op, argument
876 if hasattr(method, "_sa_instrument_after"):
877 op = method._sa_instrument_after
878 assert op in ("fire_append_event", "fire_remove_event")
879 after = op
880 if before:
881 methods[name] = before + (after,)
882 elif after:
883 methods[name] = None, None, after
884 return roles, methods
887def _setup_canned_roles(cls, roles, methods):
888 """see if this class has "canned" roles based on a known
889 collection type (dict, set, list). Apply those roles
890 as needed to the "roles" dictionary, and also
891 prepare "decorator" methods
893 """
894 collection_type = util.duck_type_collection(cls)
895 if collection_type in __interfaces:
896 assert collection_type is not None
897 canned_roles, decorators = __interfaces[collection_type]
898 for role, name in canned_roles.items():
899 roles.setdefault(role, name)
901 # apply ABC auto-decoration to methods that need it
902 for method, decorator in decorators.items():
903 fn = getattr(cls, method, None)
904 if (
905 fn
906 and method not in methods
907 and not hasattr(fn, "_sa_instrumented")
908 ):
909 setattr(cls, method, decorator(fn))
912def _assert_required_roles(cls, roles, methods):
913 """ensure all roles are present, and apply implicit instrumentation if
914 needed
916 """
917 if "appender" not in roles or not hasattr(cls, roles["appender"]):
918 raise sa_exc.ArgumentError(
919 "Type %s must elect an appender method to be "
920 "a collection class" % cls.__name__
921 )
922 elif roles["appender"] not in methods and not hasattr(
923 getattr(cls, roles["appender"]), "_sa_instrumented"
924 ):
925 methods[roles["appender"]] = ("fire_append_event", 1, None)
927 if "remover" not in roles or not hasattr(cls, roles["remover"]):
928 raise sa_exc.ArgumentError(
929 "Type %s must elect a remover method to be "
930 "a collection class" % cls.__name__
931 )
932 elif roles["remover"] not in methods and not hasattr(
933 getattr(cls, roles["remover"]), "_sa_instrumented"
934 ):
935 methods[roles["remover"]] = ("fire_remove_event", 1, None)
937 if "iterator" not in roles or not hasattr(cls, roles["iterator"]):
938 raise sa_exc.ArgumentError(
939 "Type %s must elect an iterator method to be "
940 "a collection class" % cls.__name__
941 )
944def _set_collection_attributes(cls, roles, methods):
945 """apply ad-hoc instrumentation from decorators, class-level defaults
946 and implicit role declarations
948 """
949 for method_name, (before, argument, after) in methods.items():
950 setattr(
951 cls,
952 method_name,
953 _instrument_membership_mutator(
954 getattr(cls, method_name), before, argument, after
955 ),
956 )
957 # intern the role map
958 for role, method_name in roles.items():
959 setattr(cls, "_sa_%s" % role, getattr(cls, method_name))
961 cls._sa_adapter = None
963 cls._sa_instrumented = id(cls)
966def _instrument_membership_mutator(method, before, argument, after):
967 """Route method args and/or return value through the collection
968 adapter."""
969 # This isn't smart enough to handle @adds(1) for 'def fn(self, (a, b))'
970 if before:
971 fn_args = list(
972 util.flatten_iterator(inspect_getfullargspec(method)[0])
973 )
974 if isinstance(argument, int):
975 pos_arg = argument
976 named_arg = len(fn_args) > argument and fn_args[argument] or None
977 else:
978 if argument in fn_args:
979 pos_arg = fn_args.index(argument)
980 else:
981 pos_arg = None
982 named_arg = argument
983 del fn_args
985 def wrapper(*args, **kw):
986 if before:
987 if pos_arg is None:
988 if named_arg not in kw:
989 raise sa_exc.ArgumentError(
990 "Missing argument %s" % argument
991 )
992 value = kw[named_arg]
993 else:
994 if len(args) > pos_arg:
995 value = args[pos_arg]
996 elif named_arg in kw:
997 value = kw[named_arg]
998 else:
999 raise sa_exc.ArgumentError(
1000 "Missing argument %s" % argument
1001 )
1003 initiator = kw.pop("_sa_initiator", None)
1004 if initiator is False:
1005 executor = None
1006 else:
1007 executor = args[0]._sa_adapter
1009 if before and executor:
1010 getattr(executor, before)(value, initiator)
1012 if not after or not executor:
1013 return method(*args, **kw)
1014 else:
1015 res = method(*args, **kw)
1016 if res is not None:
1017 getattr(executor, after)(res, initiator)
1018 return res
1020 wrapper._sa_instrumented = True # type: ignore[attr-defined]
1021 if hasattr(method, "_sa_instrument_role"):
1022 wrapper._sa_instrument_role = method._sa_instrument_role # type: ignore[attr-defined] # noqa: E501
1023 wrapper.__name__ = method.__name__
1024 wrapper.__doc__ = method.__doc__
1025 return wrapper
1028def __set_wo_mutation(collection, item, _sa_initiator=None):
1029 """Run set wo mutation events.
1031 The collection is not mutated.
1033 """
1034 if _sa_initiator is not False:
1035 executor = collection._sa_adapter
1036 if executor:
1037 executor.fire_append_wo_mutation_event(
1038 item, _sa_initiator, key=None
1039 )
1042def __set(collection, item, _sa_initiator, key):
1043 """Run set events.
1045 This event always occurs before the collection is actually mutated.
1047 """
1049 if _sa_initiator is not False:
1050 executor = collection._sa_adapter
1051 if executor:
1052 item = executor.fire_append_event(item, _sa_initiator, key=key)
1053 return item
1056def __del(collection, item, _sa_initiator, key):
1057 """Run del events.
1059 This event occurs before the collection is actually mutated, *except*
1060 in the case of a pop operation, in which case it occurs afterwards.
1061 For pop operations, the __before_pop hook is called before the
1062 operation occurs.
1064 """
1065 if _sa_initiator is not False:
1066 executor = collection._sa_adapter
1067 if executor:
1068 executor.fire_remove_event(item, _sa_initiator, key=key)
1071def __before_pop(collection, _sa_initiator=None):
1072 """An event which occurs on a before a pop() operation occurs."""
1073 executor = collection._sa_adapter
1074 if executor:
1075 executor.fire_pre_remove_event(_sa_initiator)
1078def _list_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1079 """Tailored instrumentation wrappers for any list-like class."""
1081 def _tidy(fn):
1082 fn._sa_instrumented = True
1083 fn.__doc__ = getattr(list, fn.__name__).__doc__
1085 def append(fn):
1086 def append(self, item, _sa_initiator=None):
1087 item = __set(self, item, _sa_initiator, NO_KEY)
1088 fn(self, item)
1090 _tidy(append)
1091 return append
1093 def remove(fn):
1094 def remove(self, value, _sa_initiator=None):
1095 __del(self, value, _sa_initiator, NO_KEY)
1096 # testlib.pragma exempt:__eq__
1097 fn(self, value)
1099 _tidy(remove)
1100 return remove
1102 def insert(fn):
1103 def insert(self, index, value):
1104 value = __set(self, value, None, index)
1105 fn(self, index, value)
1107 _tidy(insert)
1108 return insert
1110 def __setitem__(fn):
1111 def __setitem__(self, index, value):
1112 if not isinstance(index, slice):
1113 existing = self[index]
1114 if existing is not None:
1115 __del(self, existing, None, index)
1116 value = __set(self, value, None, index)
1117 fn(self, index, value)
1118 else:
1119 # slice assignment requires __delitem__, insert, __len__
1120 step = index.step or 1
1121 start = index.start or 0
1122 if start < 0:
1123 start += len(self)
1124 if index.stop is not None:
1125 stop = index.stop
1126 else:
1127 stop = len(self)
1128 if stop < 0:
1129 stop += len(self)
1131 if step == 1:
1132 if value is self:
1133 return
1134 for i in range(start, stop, step):
1135 if len(self) > start:
1136 del self[start]
1138 for i, item in enumerate(value):
1139 self.insert(i + start, item)
1140 else:
1141 rng = list(range(start, stop, step))
1142 if len(value) != len(rng):
1143 raise ValueError(
1144 "attempt to assign sequence of size %s to "
1145 "extended slice of size %s"
1146 % (len(value), len(rng))
1147 )
1148 for i, item in zip(rng, value):
1149 self.__setitem__(i, item)
1151 _tidy(__setitem__)
1152 return __setitem__
1154 def __delitem__(fn):
1155 def __delitem__(self, index):
1156 if not isinstance(index, slice):
1157 item = self[index]
1158 __del(self, item, None, index)
1159 fn(self, index)
1160 else:
1161 # slice deletion requires __getslice__ and a slice-groking
1162 # __getitem__ for stepped deletion
1163 # note: not breaking this into atomic dels
1164 for item in self[index]:
1165 __del(self, item, None, index)
1166 fn(self, index)
1168 _tidy(__delitem__)
1169 return __delitem__
1171 def extend(fn):
1172 def extend(self, iterable):
1173 for value in list(iterable):
1174 self.append(value)
1176 _tidy(extend)
1177 return extend
1179 def __iadd__(fn):
1180 def __iadd__(self, iterable):
1181 # list.__iadd__ takes any iterable and seems to let TypeError
1182 # raise as-is instead of returning NotImplemented
1183 for value in list(iterable):
1184 self.append(value)
1185 return self
1187 _tidy(__iadd__)
1188 return __iadd__
1190 def pop(fn):
1191 def pop(self, index=-1):
1192 __before_pop(self)
1193 item = fn(self, index)
1194 __del(self, item, None, index)
1195 return item
1197 _tidy(pop)
1198 return pop
1200 def clear(fn):
1201 def clear(self, index=-1):
1202 for item in self:
1203 __del(self, item, None, index)
1204 fn(self)
1206 _tidy(clear)
1207 return clear
1209 # __imul__ : not wrapping this. all members of the collection are already
1210 # present, so no need to fire appends... wrapping it with an explicit
1211 # decorator is still possible, so events on *= can be had if they're
1212 # desired. hard to imagine a use case for __imul__, though.
1214 l = locals().copy()
1215 l.pop("_tidy")
1216 return l
1219def _dict_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1220 """Tailored instrumentation wrappers for any dict-like mapping class."""
1222 def _tidy(fn):
1223 fn._sa_instrumented = True
1224 fn.__doc__ = getattr(dict, fn.__name__).__doc__
1226 def __setitem__(fn):
1227 def __setitem__(self, key, value, _sa_initiator=None):
1228 if key in self:
1229 __del(self, self[key], _sa_initiator, key)
1230 value = __set(self, value, _sa_initiator, key)
1231 fn(self, key, value)
1233 _tidy(__setitem__)
1234 return __setitem__
1236 def __delitem__(fn):
1237 def __delitem__(self, key, _sa_initiator=None):
1238 if key in self:
1239 __del(self, self[key], _sa_initiator, key)
1240 fn(self, key)
1242 _tidy(__delitem__)
1243 return __delitem__
1245 def clear(fn):
1246 def clear(self):
1247 for key in self:
1248 __del(self, self[key], None, key)
1249 fn(self)
1251 _tidy(clear)
1252 return clear
1254 def pop(fn):
1255 def pop(self, key, default=NO_ARG):
1256 __before_pop(self)
1257 _to_del = key in self
1258 if default is NO_ARG:
1259 item = fn(self, key)
1260 else:
1261 item = fn(self, key, default)
1262 if _to_del:
1263 __del(self, item, None, key)
1264 return item
1266 _tidy(pop)
1267 return pop
1269 def popitem(fn):
1270 def popitem(self):
1271 __before_pop(self)
1272 item = fn(self)
1273 __del(self, item[1], None, 1)
1274 return item
1276 _tidy(popitem)
1277 return popitem
1279 def setdefault(fn):
1280 def setdefault(self, key, default=None):
1281 if key not in self:
1282 self.__setitem__(key, default)
1283 return default
1284 else:
1285 value = self.__getitem__(key)
1286 if value is default:
1287 __set_wo_mutation(self, value, None)
1289 return value
1291 _tidy(setdefault)
1292 return setdefault
1294 def update(fn):
1295 def update(self, __other=NO_ARG, **kw):
1296 if __other is not NO_ARG:
1297 if hasattr(__other, "keys"):
1298 for key in list(__other):
1299 if key not in self or self[key] is not __other[key]:
1300 self[key] = __other[key]
1301 else:
1302 __set_wo_mutation(self, __other[key], None)
1303 else:
1304 for key, value in __other:
1305 if key not in self or self[key] is not value:
1306 self[key] = value
1307 else:
1308 __set_wo_mutation(self, value, None)
1309 for key in kw:
1310 if key not in self or self[key] is not kw[key]:
1311 self[key] = kw[key]
1312 else:
1313 __set_wo_mutation(self, kw[key], None)
1315 _tidy(update)
1316 return update
1318 l = locals().copy()
1319 l.pop("_tidy")
1320 return l
1323_set_binop_bases = (set, frozenset)
1326def _set_binops_check_strict(self: Any, obj: Any) -> bool:
1327 """Allow only set, frozenset and self.__class__-derived
1328 objects in binops."""
1329 return isinstance(obj, _set_binop_bases + (self.__class__,))
1332def _set_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1333 """Tailored instrumentation wrappers for any set-like class."""
1335 def _tidy(fn):
1336 fn._sa_instrumented = True
1337 fn.__doc__ = getattr(set, fn.__name__).__doc__
1339 def add(fn):
1340 def add(self, value, _sa_initiator=None):
1341 if value not in self:
1342 value = __set(self, value, _sa_initiator, NO_KEY)
1343 else:
1344 __set_wo_mutation(self, value, _sa_initiator)
1345 # testlib.pragma exempt:__hash__
1346 fn(self, value)
1348 _tidy(add)
1349 return add
1351 def discard(fn):
1352 def discard(self, value, _sa_initiator=None):
1353 # testlib.pragma exempt:__hash__
1354 if value in self:
1355 __del(self, value, _sa_initiator, NO_KEY)
1356 # testlib.pragma exempt:__hash__
1357 fn(self, value)
1359 _tidy(discard)
1360 return discard
1362 def remove(fn):
1363 def remove(self, value, _sa_initiator=None):
1364 # testlib.pragma exempt:__hash__
1365 if value in self:
1366 __del(self, value, _sa_initiator, NO_KEY)
1367 # testlib.pragma exempt:__hash__
1368 fn(self, value)
1370 _tidy(remove)
1371 return remove
1373 def pop(fn):
1374 def pop(self):
1375 __before_pop(self)
1376 item = fn(self)
1377 # for set in particular, we have no way to access the item
1378 # that will be popped before pop is called.
1379 __del(self, item, None, NO_KEY)
1380 return item
1382 _tidy(pop)
1383 return pop
1385 def clear(fn):
1386 def clear(self):
1387 for item in list(self):
1388 self.remove(item)
1390 _tidy(clear)
1391 return clear
1393 def update(fn):
1394 def update(self, value):
1395 for item in value:
1396 self.add(item)
1398 _tidy(update)
1399 return update
1401 def __ior__(fn):
1402 def __ior__(self, value):
1403 if not _set_binops_check_strict(self, value):
1404 return NotImplemented
1405 for item in value:
1406 self.add(item)
1407 return self
1409 _tidy(__ior__)
1410 return __ior__
1412 def difference_update(fn):
1413 def difference_update(self, value):
1414 for item in value:
1415 self.discard(item)
1417 _tidy(difference_update)
1418 return difference_update
1420 def __isub__(fn):
1421 def __isub__(self, value):
1422 if not _set_binops_check_strict(self, value):
1423 return NotImplemented
1424 for item in value:
1425 self.discard(item)
1426 return self
1428 _tidy(__isub__)
1429 return __isub__
1431 def intersection_update(fn):
1432 def intersection_update(self, other):
1433 want, have = self.intersection(other), set(self)
1434 remove, add = have - want, want - have
1436 for item in remove:
1437 self.remove(item)
1438 for item in add:
1439 self.add(item)
1441 _tidy(intersection_update)
1442 return intersection_update
1444 def __iand__(fn):
1445 def __iand__(self, other):
1446 if not _set_binops_check_strict(self, other):
1447 return NotImplemented
1448 want, have = self.intersection(other), set(self)
1449 remove, add = have - want, want - have
1451 for item in remove:
1452 self.remove(item)
1453 for item in add:
1454 self.add(item)
1455 return self
1457 _tidy(__iand__)
1458 return __iand__
1460 def symmetric_difference_update(fn):
1461 def symmetric_difference_update(self, other):
1462 want, have = self.symmetric_difference(other), set(self)
1463 remove, add = have - want, want - have
1465 for item in remove:
1466 self.remove(item)
1467 for item in add:
1468 self.add(item)
1470 _tidy(symmetric_difference_update)
1471 return symmetric_difference_update
1473 def __ixor__(fn):
1474 def __ixor__(self, other):
1475 if not _set_binops_check_strict(self, other):
1476 return NotImplemented
1477 want, have = self.symmetric_difference(other), set(self)
1478 remove, add = have - want, want - have
1480 for item in remove:
1481 self.remove(item)
1482 for item in add:
1483 self.add(item)
1484 return self
1486 _tidy(__ixor__)
1487 return __ixor__
1489 l = locals().copy()
1490 l.pop("_tidy")
1491 return l
1494class InstrumentedList(List[_T]):
1495 """An instrumented version of the built-in list."""
1498class InstrumentedSet(Set[_T]):
1499 """An instrumented version of the built-in set."""
1502class InstrumentedDict(Dict[_KT, _VT]):
1503 """An instrumented version of the built-in dict."""
1506__canned_instrumentation = cast(
1507 util.immutabledict[Any, _CollectionFactoryType],
1508 util.immutabledict(
1509 {
1510 list: InstrumentedList,
1511 set: InstrumentedSet,
1512 dict: InstrumentedDict,
1513 }
1514 ),
1515)
1517__interfaces: util.immutabledict[
1518 Any,
1519 Tuple[
1520 Dict[str, str],
1521 Dict[str, Callable[..., Any]],
1522 ],
1523] = util.immutabledict(
1524 {
1525 list: (
1526 {
1527 "appender": "append",
1528 "remover": "remove",
1529 "iterator": "__iter__",
1530 },
1531 _list_decorators(),
1532 ),
1533 set: (
1534 {"appender": "add", "remover": "remove", "iterator": "__iter__"},
1535 _set_decorators(),
1536 ),
1537 # decorators are required for dicts and object collections.
1538 dict: ({"iterator": "values"}, _dict_decorators()),
1539 }
1540)
1543def __go(lcls):
1544 global keyfunc_mapping, mapped_collection
1545 global column_keyed_dict, column_mapped_collection
1546 global MappedCollection, KeyFuncDict
1547 global attribute_keyed_dict, attribute_mapped_collection
1549 from .mapped_collection import keyfunc_mapping
1550 from .mapped_collection import column_keyed_dict
1551 from .mapped_collection import attribute_keyed_dict
1552 from .mapped_collection import KeyFuncDict
1554 from .mapped_collection import mapped_collection
1555 from .mapped_collection import column_mapped_collection
1556 from .mapped_collection import attribute_mapped_collection
1557 from .mapped_collection import MappedCollection
1559 # ensure instrumentation is associated with
1560 # these built-in classes; if a user-defined class
1561 # subclasses these and uses @internally_instrumented,
1562 # the superclass is otherwise not instrumented.
1563 # see [ticket:2406].
1564 _instrument_class(InstrumentedList)
1565 _instrument_class(InstrumentedSet)
1566 _instrument_class(KeyFuncDict)
1569__go(locals())