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 Set
124from typing import Tuple
125from typing import Type
126from typing import TYPE_CHECKING
127from typing import TypeVar
128from typing import Union
129import weakref
131from .base import NO_KEY
132from .. import exc as sa_exc
133from .. import util
134from ..sql.base import NO_ARG
135from ..util.compat import inspect_getfullargspec
136from ..util.typing import Protocol
138if typing.TYPE_CHECKING:
139 from .attributes import AttributeEventToken
140 from .attributes import CollectionAttributeImpl
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]]
183 _sa_converter: _CollectionConverterProtocol
186class collection:
187 """Decorators for entity collection classes.
189 The decorators fall into two groups: annotations and interception recipes.
191 The annotating decorators (appender, remover, iterator, converter,
192 internally_instrumented) indicate the method's purpose and take no
193 arguments. They are not written with parens::
195 @collection.appender
196 def append(self, append): ...
198 The recipe decorators all require parens, even those that take no
199 arguments::
201 @collection.adds("entity")
202 def insert(self, position, entity): ...
205 @collection.removes_return()
206 def popitem(self): ...
208 """
210 # Bundled as a class solely for ease of use: packaging, doc strings,
211 # importability.
213 @staticmethod
214 def appender(fn):
215 """Tag the method as the collection appender.
217 The appender method is called with one positional argument: the value
218 to append. The method will be automatically decorated with 'adds(1)'
219 if not already decorated::
221 @collection.appender
222 def add(self, append): ...
225 # or, equivalently
226 @collection.appender
227 @collection.adds(1)
228 def add(self, append): ...
231 # for mapping type, an 'append' may kick out a previous value
232 # that occupies that slot. consider d['a'] = 'foo'- any previous
233 # value in d['a'] is discarded.
234 @collection.appender
235 @collection.replaces(1)
236 def add(self, entity):
237 key = some_key_func(entity)
238 previous = None
239 if key in self:
240 previous = self[key]
241 self[key] = entity
242 return previous
244 If the value to append is not allowed in the collection, you may
245 raise an exception. Something to remember is that the appender
246 will be called for each object mapped by a database query. If the
247 database contains rows that violate your collection semantics, you
248 will need to get creative to fix the problem, as access via the
249 collection will not work.
251 If the appender method is internally instrumented, you must also
252 receive the keyword argument '_sa_initiator' and ensure its
253 promulgation to collection events.
255 """
256 fn._sa_instrument_role = "appender"
257 return fn
259 @staticmethod
260 def remover(fn):
261 """Tag the method as the collection remover.
263 The remover method is called with one positional argument: the value
264 to remove. The method will be automatically decorated with
265 :meth:`removes_return` if not already decorated::
267 @collection.remover
268 def zap(self, entity): ...
271 # or, equivalently
272 @collection.remover
273 @collection.removes_return()
274 def zap(self): ...
276 If the value to remove is not present in the collection, you may
277 raise an exception or return None to ignore the error.
279 If the remove method is internally instrumented, you must also
280 receive the keyword argument '_sa_initiator' and ensure its
281 promulgation to collection events.
283 """
284 fn._sa_instrument_role = "remover"
285 return fn
287 @staticmethod
288 def iterator(fn):
289 """Tag the method as the collection remover.
291 The iterator method is called with no arguments. It is expected to
292 return an iterator over all collection members::
294 @collection.iterator
295 def __iter__(self): ...
297 """
298 fn._sa_instrument_role = "iterator"
299 return fn
301 @staticmethod
302 def internally_instrumented(fn):
303 """Tag the method as instrumented.
305 This tag will prevent any decoration from being applied to the
306 method. Use this if you are orchestrating your own calls to
307 :func:`.collection_adapter` in one of the basic SQLAlchemy
308 interface methods, or to prevent an automatic ABC method
309 decoration from wrapping your implementation::
311 # normally an 'extend' method on a list-like class would be
312 # automatically intercepted and re-implemented in terms of
313 # SQLAlchemy events and append(). your implementation will
314 # never be called, unless:
315 @collection.internally_instrumented
316 def extend(self, items): ...
318 """
319 fn._sa_instrumented = True
320 return fn
322 @staticmethod
323 @util.deprecated(
324 "1.3",
325 "The :meth:`.collection.converter` handler is deprecated and will "
326 "be removed in a future release. Please refer to the "
327 ":class:`.AttributeEvents.bulk_replace` listener interface in "
328 "conjunction with the :func:`.event.listen` function.",
329 )
330 def converter(fn):
331 """Tag the method as the collection converter.
333 This optional method will be called when a collection is being
334 replaced entirely, as in::
336 myobj.acollection = [newvalue1, newvalue2]
338 The converter method will receive the object being assigned and should
339 return an iterable of values suitable for use by the ``appender``
340 method. A converter must not assign values or mutate the collection,
341 its sole job is to adapt the value the user provides into an iterable
342 of values for the ORM's use.
344 The default converter implementation will use duck-typing to do the
345 conversion. A dict-like collection will be convert into an iterable
346 of dictionary values, and other types will simply be iterated::
348 @collection.converter
349 def convert(self, other): ...
351 If the duck-typing of the object does not match the type of this
352 collection, a TypeError is raised.
354 Supply an implementation of this method if you want to expand the
355 range of possible types that can be assigned in bulk or perform
356 validation on the values about to be assigned.
358 """
359 fn._sa_instrument_role = "converter"
360 return fn
362 @staticmethod
363 def adds(arg: int) -> Callable[[_FN], _FN]:
364 """Mark the method as adding an entity to the collection.
366 Adds "add to collection" handling to the method. The decorator
367 argument indicates which method argument holds the SQLAlchemy-relevant
368 value. Arguments can be specified positionally (i.e. integer) or by
369 name::
371 @collection.adds(1)
372 def push(self, item): ...
375 @collection.adds("entity")
376 def do_stuff(self, thing, entity=None): ...
378 """
380 def decorator(fn):
381 fn._sa_instrument_before = ("fire_append_event", arg)
382 return fn
384 return decorator
386 @staticmethod
387 def replaces(arg):
388 """Mark the method as replacing an entity in the collection.
390 Adds "add to collection" and "remove from collection" handling to
391 the method. The decorator argument indicates which method argument
392 holds the SQLAlchemy-relevant value to be added, and return value, if
393 any will be considered the value to remove.
395 Arguments can be specified positionally (i.e. integer) or by name::
397 @collection.replaces(2)
398 def __setitem__(self, index, item): ...
400 """
402 def decorator(fn):
403 fn._sa_instrument_before = ("fire_append_event", arg)
404 fn._sa_instrument_after = "fire_remove_event"
405 return fn
407 return decorator
409 @staticmethod
410 def removes(arg):
411 """Mark the method as removing an entity in the collection.
413 Adds "remove from collection" handling to the method. The decorator
414 argument indicates which method argument holds the SQLAlchemy-relevant
415 value to be removed. Arguments can be specified positionally (i.e.
416 integer) or by name::
418 @collection.removes(1)
419 def zap(self, item): ...
421 For methods where the value to remove is not known at call-time, use
422 collection.removes_return.
424 """
426 def decorator(fn):
427 fn._sa_instrument_before = ("fire_remove_event", arg)
428 return fn
430 return decorator
432 @staticmethod
433 def removes_return():
434 """Mark the method as removing an entity in the collection.
436 Adds "remove from collection" handling to the method. The return
437 value of the method, if any, is considered the value to remove. The
438 method arguments are not inspected::
440 @collection.removes_return()
441 def pop(self): ...
443 For methods where the value to remove is known at call-time, use
444 collection.remove.
446 """
448 def decorator(fn):
449 fn._sa_instrument_after = "fire_remove_event"
450 return fn
452 return decorator
455if TYPE_CHECKING:
457 def collection_adapter(collection: Collection[Any]) -> CollectionAdapter:
458 """Fetch the :class:`.CollectionAdapter` for a collection."""
460else:
461 collection_adapter = operator.attrgetter("_sa_adapter")
464class CollectionAdapter:
465 """Bridges between the ORM and arbitrary Python collections.
467 Proxies base-level collection operations (append, remove, iterate)
468 to the underlying Python collection, and emits add/remove events for
469 entities entering or leaving the collection.
471 The ORM uses :class:`.CollectionAdapter` exclusively for interaction with
472 entity collections.
475 """
477 __slots__ = (
478 "attr",
479 "_key",
480 "_data",
481 "owner_state",
482 "_converter",
483 "invalidated",
484 "empty",
485 )
487 attr: CollectionAttributeImpl
488 _key: str
490 # this is actually a weakref; see note in constructor
491 _data: Callable[..., _AdaptedCollectionProtocol]
493 owner_state: InstanceState[Any]
494 _converter: _CollectionConverterProtocol
495 invalidated: bool
496 empty: bool
498 def __init__(
499 self,
500 attr: CollectionAttributeImpl,
501 owner_state: InstanceState[Any],
502 data: _AdaptedCollectionProtocol,
503 ):
504 self.attr = attr
505 self._key = attr.key
507 # this weakref stays referenced throughout the lifespan of
508 # CollectionAdapter. so while the weakref can return None, this
509 # is realistically only during garbage collection of this object, so
510 # we type this as a callable that returns _AdaptedCollectionProtocol
511 # in all cases.
512 self._data = weakref.ref(data) # type: ignore
514 self.owner_state = owner_state
515 data._sa_adapter = self
516 self._converter = data._sa_converter
517 self.invalidated = False
518 self.empty = False
520 def _warn_invalidated(self) -> None:
521 util.warn("This collection has been invalidated.")
523 @property
524 def data(self) -> _AdaptedCollectionProtocol:
525 "The entity collection being adapted."
526 return self._data()
528 @property
529 def _referenced_by_owner(self) -> bool:
530 """return True if the owner state still refers to this collection.
532 This will return False within a bulk replace operation,
533 where this collection is the one being replaced.
535 """
536 return self.owner_state.dict[self._key] is self._data()
538 def bulk_appender(self):
539 return self._data()._sa_appender
541 def append_with_event(
542 self, item: Any, initiator: Optional[AttributeEventToken] = None
543 ) -> None:
544 """Add an entity to the collection, firing mutation events."""
546 self._data()._sa_appender(item, _sa_initiator=initiator)
548 def _set_empty(self, user_data):
549 assert (
550 not self.empty
551 ), "This collection adapter is already in the 'empty' state"
552 self.empty = True
553 self.owner_state._empty_collections[self._key] = user_data
555 def _reset_empty(self) -> None:
556 assert (
557 self.empty
558 ), "This collection adapter is not in the 'empty' state"
559 self.empty = False
560 self.owner_state.dict[self._key] = (
561 self.owner_state._empty_collections.pop(self._key)
562 )
564 def _refuse_empty(self) -> NoReturn:
565 raise sa_exc.InvalidRequestError(
566 "This is a special 'empty' collection which cannot accommodate "
567 "internal mutation operations"
568 )
570 def append_without_event(self, item: Any) -> None:
571 """Add or restore an entity to the collection, firing no events."""
573 if self.empty:
574 self._refuse_empty()
575 self._data()._sa_appender(item, _sa_initiator=False)
577 def append_multiple_without_event(self, items: Iterable[Any]) -> None:
578 """Add or restore an entity to the collection, firing no events."""
579 if self.empty:
580 self._refuse_empty()
581 appender = self._data()._sa_appender
582 for item in items:
583 appender(item, _sa_initiator=False)
585 def bulk_remover(self):
586 return self._data()._sa_remover
588 def remove_with_event(
589 self, item: Any, initiator: Optional[AttributeEventToken] = None
590 ) -> None:
591 """Remove an entity from the collection, firing mutation events."""
592 self._data()._sa_remover(item, _sa_initiator=initiator)
594 def remove_without_event(self, item: Any) -> None:
595 """Remove an entity from the collection, firing no events."""
596 if self.empty:
597 self._refuse_empty()
598 self._data()._sa_remover(item, _sa_initiator=False)
600 def clear_with_event(
601 self, initiator: Optional[AttributeEventToken] = None
602 ) -> None:
603 """Empty the collection, firing a mutation event for each entity."""
605 if self.empty:
606 self._refuse_empty()
607 remover = self._data()._sa_remover
608 for item in list(self):
609 remover(item, _sa_initiator=initiator)
611 def clear_without_event(self) -> None:
612 """Empty the collection, firing no events."""
614 if self.empty:
615 self._refuse_empty()
616 remover = self._data()._sa_remover
617 for item in list(self):
618 remover(item, _sa_initiator=False)
620 def __iter__(self):
621 """Iterate over entities in the collection."""
623 return iter(self._data()._sa_iterator())
625 def __len__(self):
626 """Count entities in the collection."""
627 return len(list(self._data()._sa_iterator()))
629 def __bool__(self):
630 return True
632 def _fire_append_wo_mutation_event_bulk(
633 self, items, initiator=None, key=NO_KEY
634 ):
635 if not items:
636 return
638 if initiator is not False:
639 if self.invalidated:
640 self._warn_invalidated()
642 if self.empty:
643 self._reset_empty()
645 for item in items:
646 self.attr.fire_append_wo_mutation_event(
647 self.owner_state,
648 self.owner_state.dict,
649 item,
650 initiator,
651 key,
652 )
654 def fire_append_wo_mutation_event(self, item, initiator=None, key=NO_KEY):
655 """Notify that a entity is entering the collection but is already
656 present.
659 Initiator is a token owned by the InstrumentedAttribute that
660 initiated the membership mutation, and should be left as None
661 unless you are passing along an initiator value from a chained
662 operation.
664 .. versionadded:: 1.4.15
666 """
667 if initiator is not False:
668 if self.invalidated:
669 self._warn_invalidated()
671 if self.empty:
672 self._reset_empty()
674 return self.attr.fire_append_wo_mutation_event(
675 self.owner_state, self.owner_state.dict, item, initiator, key
676 )
677 else:
678 return item
680 def fire_append_event(self, item, initiator=None, key=NO_KEY):
681 """Notify that a entity has entered the collection.
683 Initiator is a token owned by the InstrumentedAttribute that
684 initiated the membership mutation, and should be left as None
685 unless you are passing along an initiator value from a chained
686 operation.
688 """
689 if initiator is not False:
690 if self.invalidated:
691 self._warn_invalidated()
693 if self.empty:
694 self._reset_empty()
696 return self.attr.fire_append_event(
697 self.owner_state, self.owner_state.dict, item, initiator, key
698 )
699 else:
700 return item
702 def _fire_remove_event_bulk(self, items, initiator=None, key=NO_KEY):
703 if not items:
704 return
706 if initiator is not False:
707 if self.invalidated:
708 self._warn_invalidated()
710 if self.empty:
711 self._reset_empty()
713 for item in items:
714 self.attr.fire_remove_event(
715 self.owner_state,
716 self.owner_state.dict,
717 item,
718 initiator,
719 key,
720 )
722 def fire_remove_event(self, item, initiator=None, key=NO_KEY):
723 """Notify that a entity has been removed from the collection.
725 Initiator is the InstrumentedAttribute that initiated the membership
726 mutation, and should be left as None unless you are passing along
727 an initiator value from a chained operation.
729 """
730 if initiator is not False:
731 if self.invalidated:
732 self._warn_invalidated()
734 if self.empty:
735 self._reset_empty()
737 self.attr.fire_remove_event(
738 self.owner_state, self.owner_state.dict, item, initiator, key
739 )
741 def fire_pre_remove_event(self, initiator=None, key=NO_KEY):
742 """Notify that an entity is about to be removed from the collection.
744 Only called if the entity cannot be removed after calling
745 fire_remove_event().
747 """
748 if self.invalidated:
749 self._warn_invalidated()
750 self.attr.fire_pre_remove_event(
751 self.owner_state,
752 self.owner_state.dict,
753 initiator=initiator,
754 key=key,
755 )
757 def __getstate__(self):
758 return {
759 "key": self._key,
760 "owner_state": self.owner_state,
761 "owner_cls": self.owner_state.class_,
762 "data": self.data,
763 "invalidated": self.invalidated,
764 "empty": self.empty,
765 }
767 def __setstate__(self, d):
768 self._key = d["key"]
769 self.owner_state = d["owner_state"]
771 # see note in constructor regarding this type: ignore
772 self._data = weakref.ref(d["data"]) # type: ignore
774 self._converter = d["data"]._sa_converter
775 d["data"]._sa_adapter = self
776 self.invalidated = d["invalidated"]
777 self.attr = getattr(d["owner_cls"], self._key).impl
778 self.empty = d.get("empty", False)
781def bulk_replace(values, existing_adapter, new_adapter, initiator=None):
782 """Load a new collection, firing events based on prior like membership.
784 Appends instances in ``values`` onto the ``new_adapter``. Events will be
785 fired for any instance not present in the ``existing_adapter``. Any
786 instances in ``existing_adapter`` not present in ``values`` will have
787 remove events fired upon them.
789 :param values: An iterable of collection member instances
791 :param existing_adapter: A :class:`.CollectionAdapter` of
792 instances to be replaced
794 :param new_adapter: An empty :class:`.CollectionAdapter`
795 to load with ``values``
798 """
800 assert isinstance(values, list)
802 idset = util.IdentitySet
803 existing_idset = idset(existing_adapter or ())
804 constants = existing_idset.intersection(values or ())
805 additions = idset(values or ()).difference(constants)
806 removals = existing_idset.difference(constants)
808 appender = new_adapter.bulk_appender()
810 for member in values or ():
811 if member in additions:
812 appender(member, _sa_initiator=initiator)
813 elif member in constants:
814 appender(member, _sa_initiator=False)
816 if existing_adapter:
817 existing_adapter._fire_append_wo_mutation_event_bulk(
818 constants, initiator=initiator
819 )
820 existing_adapter._fire_remove_event_bulk(removals, initiator=initiator)
823def prepare_instrumentation(
824 factory: Union[Type[Collection[Any]], _CollectionFactoryType],
825) -> _CollectionFactoryType:
826 """Prepare a callable for future use as a collection class factory.
828 Given a collection class factory (either a type or no-arg callable),
829 return another factory that will produce compatible instances when
830 called.
832 This function is responsible for converting collection_class=list
833 into the run-time behavior of collection_class=InstrumentedList.
835 """
837 impl_factory: _CollectionFactoryType
839 # Convert a builtin to 'Instrumented*'
840 if factory in __canned_instrumentation:
841 impl_factory = __canned_instrumentation[factory]
842 else:
843 impl_factory = cast(_CollectionFactoryType, factory)
845 cls: Union[_CollectionFactoryType, Type[Collection[Any]]]
847 # Create a specimen
848 cls = type(impl_factory())
850 # Did factory callable return a builtin?
851 if cls in __canned_instrumentation:
852 # if so, just convert.
853 # in previous major releases, this codepath wasn't working and was
854 # not covered by tests. prior to that it supplied a "wrapper"
855 # function that would return the class, though the rationale for this
856 # case is not known
857 impl_factory = __canned_instrumentation[cls]
858 cls = type(impl_factory())
860 # Instrument the class if needed.
861 if __instrumentation_mutex.acquire():
862 try:
863 if getattr(cls, "_sa_instrumented", None) != id(cls):
864 _instrument_class(cls)
865 finally:
866 __instrumentation_mutex.release()
868 return impl_factory
871def _instrument_class(cls):
872 """Modify methods in a class and install instrumentation."""
874 # In the normal call flow, a request for any of the 3 basic collection
875 # types is transformed into one of our trivial subclasses
876 # (e.g. InstrumentedList). Catch anything else that sneaks in here...
877 if cls.__module__ == "__builtin__":
878 raise sa_exc.ArgumentError(
879 "Can not instrument a built-in type. Use a "
880 "subclass, even a trivial one."
881 )
883 roles, methods = _locate_roles_and_methods(cls)
885 _setup_canned_roles(cls, roles, methods)
887 _assert_required_roles(cls, roles, methods)
889 _set_collection_attributes(cls, roles, methods)
892def _locate_roles_and_methods(cls):
893 """search for _sa_instrument_role-decorated methods in
894 method resolution order, assign to roles.
896 """
898 roles: Dict[str, str] = {}
899 methods: Dict[str, Tuple[Optional[str], Optional[int], Optional[str]]] = {}
901 for supercls in cls.__mro__:
902 for name, method in vars(supercls).items():
903 if not callable(method):
904 continue
906 # note role declarations
907 if hasattr(method, "_sa_instrument_role"):
908 role = method._sa_instrument_role
909 assert role in (
910 "appender",
911 "remover",
912 "iterator",
913 "converter",
914 )
915 roles.setdefault(role, name)
917 # transfer instrumentation requests from decorated function
918 # to the combined queue
919 before: Optional[Tuple[str, int]] = None
920 after: Optional[str] = None
922 if hasattr(method, "_sa_instrument_before"):
923 op, argument = method._sa_instrument_before
924 assert op in ("fire_append_event", "fire_remove_event")
925 before = op, argument
926 if hasattr(method, "_sa_instrument_after"):
927 op = method._sa_instrument_after
928 assert op in ("fire_append_event", "fire_remove_event")
929 after = op
930 if before:
931 methods[name] = before + (after,)
932 elif after:
933 methods[name] = None, None, after
934 return roles, methods
937def _setup_canned_roles(cls, roles, methods):
938 """see if this class has "canned" roles based on a known
939 collection type (dict, set, list). Apply those roles
940 as needed to the "roles" dictionary, and also
941 prepare "decorator" methods
943 """
944 collection_type = util.duck_type_collection(cls)
945 if collection_type in __interfaces:
946 assert collection_type is not None
947 canned_roles, decorators = __interfaces[collection_type]
948 for role, name in canned_roles.items():
949 roles.setdefault(role, name)
951 # apply ABC auto-decoration to methods that need it
952 for method, decorator in decorators.items():
953 fn = getattr(cls, method, None)
954 if (
955 fn
956 and method not in methods
957 and not hasattr(fn, "_sa_instrumented")
958 ):
959 setattr(cls, method, decorator(fn))
962def _assert_required_roles(cls, roles, methods):
963 """ensure all roles are present, and apply implicit instrumentation if
964 needed
966 """
967 if "appender" not in roles or not hasattr(cls, roles["appender"]):
968 raise sa_exc.ArgumentError(
969 "Type %s must elect an appender method to be "
970 "a collection class" % cls.__name__
971 )
972 elif roles["appender"] not in methods and not hasattr(
973 getattr(cls, roles["appender"]), "_sa_instrumented"
974 ):
975 methods[roles["appender"]] = ("fire_append_event", 1, None)
977 if "remover" not in roles or not hasattr(cls, roles["remover"]):
978 raise sa_exc.ArgumentError(
979 "Type %s must elect a remover method to be "
980 "a collection class" % cls.__name__
981 )
982 elif roles["remover"] not in methods and not hasattr(
983 getattr(cls, roles["remover"]), "_sa_instrumented"
984 ):
985 methods[roles["remover"]] = ("fire_remove_event", 1, None)
987 if "iterator" not in roles or not hasattr(cls, roles["iterator"]):
988 raise sa_exc.ArgumentError(
989 "Type %s must elect an iterator method to be "
990 "a collection class" % cls.__name__
991 )
994def _set_collection_attributes(cls, roles, methods):
995 """apply ad-hoc instrumentation from decorators, class-level defaults
996 and implicit role declarations
998 """
999 for method_name, (before, argument, after) in methods.items():
1000 setattr(
1001 cls,
1002 method_name,
1003 _instrument_membership_mutator(
1004 getattr(cls, method_name), before, argument, after
1005 ),
1006 )
1007 # intern the role map
1008 for role, method_name in roles.items():
1009 setattr(cls, "_sa_%s" % role, getattr(cls, method_name))
1011 cls._sa_adapter = None
1013 if not hasattr(cls, "_sa_converter"):
1014 cls._sa_converter = None
1015 cls._sa_instrumented = id(cls)
1018def _instrument_membership_mutator(method, before, argument, after):
1019 """Route method args and/or return value through the collection
1020 adapter."""
1021 # This isn't smart enough to handle @adds(1) for 'def fn(self, (a, b))'
1022 if before:
1023 fn_args = list(
1024 util.flatten_iterator(inspect_getfullargspec(method)[0])
1025 )
1026 if isinstance(argument, int):
1027 pos_arg = argument
1028 named_arg = len(fn_args) > argument and fn_args[argument] or None
1029 else:
1030 if argument in fn_args:
1031 pos_arg = fn_args.index(argument)
1032 else:
1033 pos_arg = None
1034 named_arg = argument
1035 del fn_args
1037 def wrapper(*args, **kw):
1038 if before:
1039 if pos_arg is None:
1040 if named_arg not in kw:
1041 raise sa_exc.ArgumentError(
1042 "Missing argument %s" % argument
1043 )
1044 value = kw[named_arg]
1045 else:
1046 if len(args) > pos_arg:
1047 value = args[pos_arg]
1048 elif named_arg in kw:
1049 value = kw[named_arg]
1050 else:
1051 raise sa_exc.ArgumentError(
1052 "Missing argument %s" % argument
1053 )
1055 initiator = kw.pop("_sa_initiator", None)
1056 if initiator is False:
1057 executor = None
1058 else:
1059 executor = args[0]._sa_adapter
1061 if before and executor:
1062 getattr(executor, before)(value, initiator)
1064 if not after or not executor:
1065 return method(*args, **kw)
1066 else:
1067 res = method(*args, **kw)
1068 if res is not None:
1069 getattr(executor, after)(res, initiator)
1070 return res
1072 wrapper._sa_instrumented = True # type: ignore[attr-defined]
1073 if hasattr(method, "_sa_instrument_role"):
1074 wrapper._sa_instrument_role = method._sa_instrument_role # type: ignore[attr-defined] # noqa: E501
1075 wrapper.__name__ = method.__name__
1076 wrapper.__doc__ = method.__doc__
1077 return wrapper
1080def __set_wo_mutation(collection, item, _sa_initiator=None):
1081 """Run set wo mutation events.
1083 The collection is not mutated.
1085 """
1086 if _sa_initiator is not False:
1087 executor = collection._sa_adapter
1088 if executor:
1089 executor.fire_append_wo_mutation_event(
1090 item, _sa_initiator, key=None
1091 )
1094def __set(collection, item, _sa_initiator, key):
1095 """Run set events.
1097 This event always occurs before the collection is actually mutated.
1099 """
1101 if _sa_initiator is not False:
1102 executor = collection._sa_adapter
1103 if executor:
1104 item = executor.fire_append_event(item, _sa_initiator, key=key)
1105 return item
1108def __del(collection, item, _sa_initiator, key):
1109 """Run del events.
1111 This event occurs before the collection is actually mutated, *except*
1112 in the case of a pop operation, in which case it occurs afterwards.
1113 For pop operations, the __before_pop hook is called before the
1114 operation occurs.
1116 """
1117 if _sa_initiator is not False:
1118 executor = collection._sa_adapter
1119 if executor:
1120 executor.fire_remove_event(item, _sa_initiator, key=key)
1123def __before_pop(collection, _sa_initiator=None):
1124 """An event which occurs on a before a pop() operation occurs."""
1125 executor = collection._sa_adapter
1126 if executor:
1127 executor.fire_pre_remove_event(_sa_initiator)
1130def _list_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1131 """Tailored instrumentation wrappers for any list-like class."""
1133 def _tidy(fn):
1134 fn._sa_instrumented = True
1135 fn.__doc__ = getattr(list, fn.__name__).__doc__
1137 def append(fn):
1138 def append(self, item, _sa_initiator=None):
1139 item = __set(self, item, _sa_initiator, NO_KEY)
1140 fn(self, item)
1142 _tidy(append)
1143 return append
1145 def remove(fn):
1146 def remove(self, value, _sa_initiator=None):
1147 __del(self, value, _sa_initiator, NO_KEY)
1148 # testlib.pragma exempt:__eq__
1149 fn(self, value)
1151 _tidy(remove)
1152 return remove
1154 def insert(fn):
1155 def insert(self, index, value):
1156 value = __set(self, value, None, index)
1157 fn(self, index, value)
1159 _tidy(insert)
1160 return insert
1162 def __setitem__(fn):
1163 def __setitem__(self, index, value):
1164 if not isinstance(index, slice):
1165 existing = self[index]
1166 if existing is not None:
1167 __del(self, existing, None, index)
1168 value = __set(self, value, None, index)
1169 fn(self, index, value)
1170 else:
1171 # slice assignment requires __delitem__, insert, __len__
1172 step = index.step or 1
1173 start = index.start or 0
1174 if start < 0:
1175 start += len(self)
1176 if index.stop is not None:
1177 stop = index.stop
1178 else:
1179 stop = len(self)
1180 if stop < 0:
1181 stop += len(self)
1183 if step == 1:
1184 if value is self:
1185 return
1186 for i in range(start, stop, step):
1187 if len(self) > start:
1188 del self[start]
1190 for i, item in enumerate(value):
1191 self.insert(i + start, item)
1192 else:
1193 rng = list(range(start, stop, step))
1194 if len(value) != len(rng):
1195 raise ValueError(
1196 "attempt to assign sequence of size %s to "
1197 "extended slice of size %s"
1198 % (len(value), len(rng))
1199 )
1200 for i, item in zip(rng, value):
1201 self.__setitem__(i, item)
1203 _tidy(__setitem__)
1204 return __setitem__
1206 def __delitem__(fn):
1207 def __delitem__(self, index):
1208 if not isinstance(index, slice):
1209 item = self[index]
1210 __del(self, item, None, index)
1211 fn(self, index)
1212 else:
1213 # slice deletion requires __getslice__ and a slice-groking
1214 # __getitem__ for stepped deletion
1215 # note: not breaking this into atomic dels
1216 for item in self[index]:
1217 __del(self, item, None, index)
1218 fn(self, index)
1220 _tidy(__delitem__)
1221 return __delitem__
1223 def extend(fn):
1224 def extend(self, iterable):
1225 for value in list(iterable):
1226 self.append(value)
1228 _tidy(extend)
1229 return extend
1231 def __iadd__(fn):
1232 def __iadd__(self, iterable):
1233 # list.__iadd__ takes any iterable and seems to let TypeError
1234 # raise as-is instead of returning NotImplemented
1235 for value in list(iterable):
1236 self.append(value)
1237 return self
1239 _tidy(__iadd__)
1240 return __iadd__
1242 def pop(fn):
1243 def pop(self, index=-1):
1244 __before_pop(self)
1245 item = fn(self, index)
1246 __del(self, item, None, index)
1247 return item
1249 _tidy(pop)
1250 return pop
1252 def clear(fn):
1253 def clear(self, index=-1):
1254 for item in self:
1255 __del(self, item, None, index)
1256 fn(self)
1258 _tidy(clear)
1259 return clear
1261 # __imul__ : not wrapping this. all members of the collection are already
1262 # present, so no need to fire appends... wrapping it with an explicit
1263 # decorator is still possible, so events on *= can be had if they're
1264 # desired. hard to imagine a use case for __imul__, though.
1266 l = locals().copy()
1267 l.pop("_tidy")
1268 return l
1271def _dict_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1272 """Tailored instrumentation wrappers for any dict-like mapping class."""
1274 def _tidy(fn):
1275 fn._sa_instrumented = True
1276 fn.__doc__ = getattr(dict, fn.__name__).__doc__
1278 def __setitem__(fn):
1279 def __setitem__(self, key, value, _sa_initiator=None):
1280 if key in self:
1281 __del(self, self[key], _sa_initiator, key)
1282 value = __set(self, value, _sa_initiator, key)
1283 fn(self, key, value)
1285 _tidy(__setitem__)
1286 return __setitem__
1288 def __delitem__(fn):
1289 def __delitem__(self, key, _sa_initiator=None):
1290 if key in self:
1291 __del(self, self[key], _sa_initiator, key)
1292 fn(self, key)
1294 _tidy(__delitem__)
1295 return __delitem__
1297 def clear(fn):
1298 def clear(self):
1299 for key in self:
1300 __del(self, self[key], None, key)
1301 fn(self)
1303 _tidy(clear)
1304 return clear
1306 def pop(fn):
1307 def pop(self, key, default=NO_ARG):
1308 __before_pop(self)
1309 _to_del = key in self
1310 if default is NO_ARG:
1311 item = fn(self, key)
1312 else:
1313 item = fn(self, key, default)
1314 if _to_del:
1315 __del(self, item, None, key)
1316 return item
1318 _tidy(pop)
1319 return pop
1321 def popitem(fn):
1322 def popitem(self):
1323 __before_pop(self)
1324 item = fn(self)
1325 __del(self, item[1], None, 1)
1326 return item
1328 _tidy(popitem)
1329 return popitem
1331 def setdefault(fn):
1332 def setdefault(self, key, default=None):
1333 if key not in self:
1334 self.__setitem__(key, default)
1335 return default
1336 else:
1337 value = self.__getitem__(key)
1338 if value is default:
1339 __set_wo_mutation(self, value, None)
1341 return value
1343 _tidy(setdefault)
1344 return setdefault
1346 def update(fn):
1347 def update(self, __other=NO_ARG, **kw):
1348 if __other is not NO_ARG:
1349 if hasattr(__other, "keys"):
1350 for key in list(__other):
1351 if key not in self or self[key] is not __other[key]:
1352 self[key] = __other[key]
1353 else:
1354 __set_wo_mutation(self, __other[key], None)
1355 else:
1356 for key, value in __other:
1357 if key not in self or self[key] is not value:
1358 self[key] = value
1359 else:
1360 __set_wo_mutation(self, value, None)
1361 for key in kw:
1362 if key not in self or self[key] is not kw[key]:
1363 self[key] = kw[key]
1364 else:
1365 __set_wo_mutation(self, kw[key], None)
1367 _tidy(update)
1368 return update
1370 l = locals().copy()
1371 l.pop("_tidy")
1372 return l
1375_set_binop_bases = (set, frozenset)
1378def _set_binops_check_strict(self: Any, obj: Any) -> bool:
1379 """Allow only set, frozenset and self.__class__-derived
1380 objects in binops."""
1381 return isinstance(obj, _set_binop_bases + (self.__class__,))
1384def _set_binops_check_loose(self: Any, obj: Any) -> bool:
1385 """Allow anything set-like to participate in set binops."""
1386 return (
1387 isinstance(obj, _set_binop_bases + (self.__class__,))
1388 or util.duck_type_collection(obj) == set
1389 )
1392def _set_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1393 """Tailored instrumentation wrappers for any set-like class."""
1395 def _tidy(fn):
1396 fn._sa_instrumented = True
1397 fn.__doc__ = getattr(set, fn.__name__).__doc__
1399 def add(fn):
1400 def add(self, value, _sa_initiator=None):
1401 if value not in self:
1402 value = __set(self, value, _sa_initiator, NO_KEY)
1403 else:
1404 __set_wo_mutation(self, value, _sa_initiator)
1405 # testlib.pragma exempt:__hash__
1406 fn(self, value)
1408 _tidy(add)
1409 return add
1411 def discard(fn):
1412 def discard(self, value, _sa_initiator=None):
1413 # testlib.pragma exempt:__hash__
1414 if value in self:
1415 __del(self, value, _sa_initiator, NO_KEY)
1416 # testlib.pragma exempt:__hash__
1417 fn(self, value)
1419 _tidy(discard)
1420 return discard
1422 def remove(fn):
1423 def remove(self, value, _sa_initiator=None):
1424 # testlib.pragma exempt:__hash__
1425 if value in self:
1426 __del(self, value, _sa_initiator, NO_KEY)
1427 # testlib.pragma exempt:__hash__
1428 fn(self, value)
1430 _tidy(remove)
1431 return remove
1433 def pop(fn):
1434 def pop(self):
1435 __before_pop(self)
1436 item = fn(self)
1437 # for set in particular, we have no way to access the item
1438 # that will be popped before pop is called.
1439 __del(self, item, None, NO_KEY)
1440 return item
1442 _tidy(pop)
1443 return pop
1445 def clear(fn):
1446 def clear(self):
1447 for item in list(self):
1448 self.remove(item)
1450 _tidy(clear)
1451 return clear
1453 def update(fn):
1454 def update(self, value):
1455 for item in value:
1456 self.add(item)
1458 _tidy(update)
1459 return update
1461 def __ior__(fn):
1462 def __ior__(self, value):
1463 if not _set_binops_check_strict(self, value):
1464 return NotImplemented
1465 for item in value:
1466 self.add(item)
1467 return self
1469 _tidy(__ior__)
1470 return __ior__
1472 def difference_update(fn):
1473 def difference_update(self, value):
1474 for item in value:
1475 self.discard(item)
1477 _tidy(difference_update)
1478 return difference_update
1480 def __isub__(fn):
1481 def __isub__(self, value):
1482 if not _set_binops_check_strict(self, value):
1483 return NotImplemented
1484 for item in value:
1485 self.discard(item)
1486 return self
1488 _tidy(__isub__)
1489 return __isub__
1491 def intersection_update(fn):
1492 def intersection_update(self, other):
1493 want, have = self.intersection(other), set(self)
1494 remove, add = have - want, want - have
1496 for item in remove:
1497 self.remove(item)
1498 for item in add:
1499 self.add(item)
1501 _tidy(intersection_update)
1502 return intersection_update
1504 def __iand__(fn):
1505 def __iand__(self, other):
1506 if not _set_binops_check_strict(self, other):
1507 return NotImplemented
1508 want, have = self.intersection(other), set(self)
1509 remove, add = have - want, want - have
1511 for item in remove:
1512 self.remove(item)
1513 for item in add:
1514 self.add(item)
1515 return self
1517 _tidy(__iand__)
1518 return __iand__
1520 def symmetric_difference_update(fn):
1521 def symmetric_difference_update(self, other):
1522 want, have = self.symmetric_difference(other), set(self)
1523 remove, add = have - want, want - have
1525 for item in remove:
1526 self.remove(item)
1527 for item in add:
1528 self.add(item)
1530 _tidy(symmetric_difference_update)
1531 return symmetric_difference_update
1533 def __ixor__(fn):
1534 def __ixor__(self, other):
1535 if not _set_binops_check_strict(self, other):
1536 return NotImplemented
1537 want, have = self.symmetric_difference(other), set(self)
1538 remove, add = have - want, want - have
1540 for item in remove:
1541 self.remove(item)
1542 for item in add:
1543 self.add(item)
1544 return self
1546 _tidy(__ixor__)
1547 return __ixor__
1549 l = locals().copy()
1550 l.pop("_tidy")
1551 return l
1554class InstrumentedList(List[_T]):
1555 """An instrumented version of the built-in list."""
1558class InstrumentedSet(Set[_T]):
1559 """An instrumented version of the built-in set."""
1562class InstrumentedDict(Dict[_KT, _VT]):
1563 """An instrumented version of the built-in dict."""
1566__canned_instrumentation: util.immutabledict[Any, _CollectionFactoryType] = (
1567 util.immutabledict(
1568 {
1569 list: InstrumentedList,
1570 set: InstrumentedSet,
1571 dict: InstrumentedDict,
1572 }
1573 )
1574)
1576__interfaces: util.immutabledict[
1577 Any,
1578 Tuple[
1579 Dict[str, str],
1580 Dict[str, Callable[..., Any]],
1581 ],
1582] = util.immutabledict(
1583 {
1584 list: (
1585 {
1586 "appender": "append",
1587 "remover": "remove",
1588 "iterator": "__iter__",
1589 },
1590 _list_decorators(),
1591 ),
1592 set: (
1593 {"appender": "add", "remover": "remove", "iterator": "__iter__"},
1594 _set_decorators(),
1595 ),
1596 # decorators are required for dicts and object collections.
1597 dict: ({"iterator": "values"}, _dict_decorators()),
1598 }
1599)
1602def __go(lcls):
1603 global keyfunc_mapping, mapped_collection
1604 global column_keyed_dict, column_mapped_collection
1605 global MappedCollection, KeyFuncDict
1606 global attribute_keyed_dict, attribute_mapped_collection
1608 from .mapped_collection import keyfunc_mapping
1609 from .mapped_collection import column_keyed_dict
1610 from .mapped_collection import attribute_keyed_dict
1611 from .mapped_collection import KeyFuncDict
1613 from .mapped_collection import mapped_collection
1614 from .mapped_collection import column_mapped_collection
1615 from .mapped_collection import attribute_mapped_collection
1616 from .mapped_collection import MappedCollection
1618 # ensure instrumentation is associated with
1619 # these built-in classes; if a user-defined class
1620 # subclasses these and uses @internally_instrumented,
1621 # the superclass is otherwise not instrumented.
1622 # see [ticket:2406].
1623 _instrument_class(InstrumentedList)
1624 _instrument_class(InstrumentedSet)
1625 _instrument_class(KeyFuncDict)
1628__go(locals())