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-2025 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"""
108from __future__ import annotations
110import operator
111import threading
112import typing
113from typing import Any
114from typing import Callable
115from typing import cast
116from typing import Collection
117from typing import Dict
118from typing import Iterable
119from typing import List
120from typing import NoReturn
121from typing import Optional
122from typing import Set
123from typing import Tuple
124from typing import Type
125from typing import TYPE_CHECKING
126from typing import TypeVar
127from typing import Union
128import weakref
130from .base import NO_KEY
131from .. import exc as sa_exc
132from .. import util
133from ..sql.base import NO_ARG
134from ..util.compat import inspect_getfullargspec
135from ..util.typing import Protocol
137if typing.TYPE_CHECKING:
138 from .attributes import AttributeEventToken
139 from .attributes import CollectionAttributeImpl
140 from .mapped_collection import attribute_keyed_dict
141 from .mapped_collection import column_keyed_dict
142 from .mapped_collection import keyfunc_mapping
143 from .mapped_collection import KeyFuncDict # noqa: F401
144 from .state import InstanceState
147__all__ = [
148 "collection",
149 "collection_adapter",
150 "keyfunc_mapping",
151 "column_keyed_dict",
152 "attribute_keyed_dict",
153 "KeyFuncDict",
154 # old names in < 2.0
155 "mapped_collection",
156 "column_mapped_collection",
157 "attribute_mapped_collection",
158 "MappedCollection",
159]
161__instrumentation_mutex = threading.Lock()
164_CollectionFactoryType = Callable[[], "_AdaptedCollectionProtocol"]
166_T = TypeVar("_T", bound=Any)
167_KT = TypeVar("_KT", bound=Any)
168_VT = TypeVar("_VT", bound=Any)
169_COL = TypeVar("_COL", bound="Collection[Any]")
170_FN = TypeVar("_FN", bound="Callable[..., Any]")
173class _CollectionConverterProtocol(Protocol):
174 def __call__(self, collection: _COL) -> _COL: ...
177class _AdaptedCollectionProtocol(Protocol):
178 _sa_adapter: CollectionAdapter
179 _sa_appender: Callable[..., Any]
180 _sa_remover: Callable[..., Any]
181 _sa_iterator: Callable[..., Iterable[Any]]
182 _sa_converter: _CollectionConverterProtocol
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, converter,
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 @util.deprecated(
323 "1.3",
324 "The :meth:`.collection.converter` handler is deprecated and will "
325 "be removed in a future release. Please refer to the "
326 ":class:`.AttributeEvents.bulk_replace` listener interface in "
327 "conjunction with the :func:`.event.listen` function.",
328 )
329 def converter(fn):
330 """Tag the method as the collection converter.
332 This optional method will be called when a collection is being
333 replaced entirely, as in::
335 myobj.acollection = [newvalue1, newvalue2]
337 The converter method will receive the object being assigned and should
338 return an iterable of values suitable for use by the ``appender``
339 method. A converter must not assign values or mutate the collection,
340 its sole job is to adapt the value the user provides into an iterable
341 of values for the ORM's use.
343 The default converter implementation will use duck-typing to do the
344 conversion. A dict-like collection will be convert into an iterable
345 of dictionary values, and other types will simply be iterated::
347 @collection.converter
348 def convert(self, other): ...
350 If the duck-typing of the object does not match the type of this
351 collection, a TypeError is raised.
353 Supply an implementation of this method if you want to expand the
354 range of possible types that can be assigned in bulk or perform
355 validation on the values about to be assigned.
357 """
358 fn._sa_instrument_role = "converter"
359 return fn
361 @staticmethod
362 def adds(arg):
363 """Mark the method as adding an entity to the collection.
365 Adds "add to collection" handling to the method. The decorator
366 argument indicates which method argument holds the SQLAlchemy-relevant
367 value. Arguments can be specified positionally (i.e. integer) or by
368 name::
370 @collection.adds(1)
371 def push(self, item): ...
374 @collection.adds("entity")
375 def do_stuff(self, thing, entity=None): ...
377 """
379 def decorator(fn):
380 fn._sa_instrument_before = ("fire_append_event", arg)
381 return fn
383 return decorator
385 @staticmethod
386 def replaces(arg):
387 """Mark the method as replacing an entity in the collection.
389 Adds "add to collection" and "remove from collection" handling to
390 the method. The decorator argument indicates which method argument
391 holds the SQLAlchemy-relevant value to be added, and return value, if
392 any will be considered the value to remove.
394 Arguments can be specified positionally (i.e. integer) or by name::
396 @collection.replaces(2)
397 def __setitem__(self, index, item): ...
399 """
401 def decorator(fn):
402 fn._sa_instrument_before = ("fire_append_event", arg)
403 fn._sa_instrument_after = "fire_remove_event"
404 return fn
406 return decorator
408 @staticmethod
409 def removes(arg):
410 """Mark the method as removing an entity in the collection.
412 Adds "remove from collection" handling to the method. The decorator
413 argument indicates which method argument holds the SQLAlchemy-relevant
414 value to be removed. Arguments can be specified positionally (i.e.
415 integer) or by name::
417 @collection.removes(1)
418 def zap(self, item): ...
420 For methods where the value to remove is not known at call-time, use
421 collection.removes_return.
423 """
425 def decorator(fn):
426 fn._sa_instrument_before = ("fire_remove_event", arg)
427 return fn
429 return decorator
431 @staticmethod
432 def removes_return():
433 """Mark the method as removing an entity in the collection.
435 Adds "remove from collection" handling to the method. The return
436 value of the method, if any, is considered the value to remove. The
437 method arguments are not inspected::
439 @collection.removes_return()
440 def pop(self): ...
442 For methods where the value to remove is known at call-time, use
443 collection.remove.
445 """
447 def decorator(fn):
448 fn._sa_instrument_after = "fire_remove_event"
449 return fn
451 return decorator
454if TYPE_CHECKING:
456 def collection_adapter(collection: Collection[Any]) -> CollectionAdapter:
457 """Fetch the :class:`.CollectionAdapter` for a collection."""
459else:
460 collection_adapter = operator.attrgetter("_sa_adapter")
463class CollectionAdapter:
464 """Bridges between the ORM and arbitrary Python collections.
466 Proxies base-level collection operations (append, remove, iterate)
467 to the underlying Python collection, and emits add/remove events for
468 entities entering or leaving the collection.
470 The ORM uses :class:`.CollectionAdapter` exclusively for interaction with
471 entity collections.
474 """
476 __slots__ = (
477 "attr",
478 "_key",
479 "_data",
480 "owner_state",
481 "_converter",
482 "invalidated",
483 "empty",
484 )
486 attr: CollectionAttributeImpl
487 _key: str
489 # this is actually a weakref; see note in constructor
490 _data: Callable[..., _AdaptedCollectionProtocol]
492 owner_state: InstanceState[Any]
493 _converter: _CollectionConverterProtocol
494 invalidated: bool
495 empty: bool
497 def __init__(
498 self,
499 attr: CollectionAttributeImpl,
500 owner_state: InstanceState[Any],
501 data: _AdaptedCollectionProtocol,
502 ):
503 self.attr = attr
504 self._key = attr.key
506 # this weakref stays referenced throughout the lifespan of
507 # CollectionAdapter. so while the weakref can return None, this
508 # is realistically only during garbage collection of this object, so
509 # we type this as a callable that returns _AdaptedCollectionProtocol
510 # in all cases.
511 self._data = weakref.ref(data) # type: ignore
513 self.owner_state = owner_state
514 data._sa_adapter = self
515 self._converter = data._sa_converter
516 self.invalidated = False
517 self.empty = False
519 def _warn_invalidated(self) -> None:
520 util.warn("This collection has been invalidated.")
522 @property
523 def data(self) -> _AdaptedCollectionProtocol:
524 "The entity collection being adapted."
525 return self._data()
527 @property
528 def _referenced_by_owner(self) -> bool:
529 """return True if the owner state still refers to this collection.
531 This will return False within a bulk replace operation,
532 where this collection is the one being replaced.
534 """
535 return self.owner_state.dict[self._key] is self._data()
537 def bulk_appender(self):
538 return self._data()._sa_appender
540 def append_with_event(
541 self, item: Any, initiator: Optional[AttributeEventToken] = None
542 ) -> None:
543 """Add an entity to the collection, firing mutation events."""
545 self._data()._sa_appender(item, _sa_initiator=initiator)
547 def _set_empty(self, user_data):
548 assert (
549 not self.empty
550 ), "This collection adapter is already in the 'empty' state"
551 self.empty = True
552 self.owner_state._empty_collections[self._key] = user_data
554 def _reset_empty(self) -> None:
555 assert (
556 self.empty
557 ), "This collection adapter is not in the 'empty' state"
558 self.empty = False
559 self.owner_state.dict[self._key] = (
560 self.owner_state._empty_collections.pop(self._key)
561 )
563 def _refuse_empty(self) -> NoReturn:
564 raise sa_exc.InvalidRequestError(
565 "This is a special 'empty' collection which cannot accommodate "
566 "internal mutation operations"
567 )
569 def append_without_event(self, item: Any) -> None:
570 """Add or restore an entity to the collection, firing no events."""
572 if self.empty:
573 self._refuse_empty()
574 self._data()._sa_appender(item, _sa_initiator=False)
576 def append_multiple_without_event(self, items: Iterable[Any]) -> None:
577 """Add or restore an entity to the collection, firing no events."""
578 if self.empty:
579 self._refuse_empty()
580 appender = self._data()._sa_appender
581 for item in items:
582 appender(item, _sa_initiator=False)
584 def bulk_remover(self):
585 return self._data()._sa_remover
587 def remove_with_event(
588 self, item: Any, initiator: Optional[AttributeEventToken] = None
589 ) -> None:
590 """Remove an entity from the collection, firing mutation events."""
591 self._data()._sa_remover(item, _sa_initiator=initiator)
593 def remove_without_event(self, item: Any) -> None:
594 """Remove an entity from the collection, firing no events."""
595 if self.empty:
596 self._refuse_empty()
597 self._data()._sa_remover(item, _sa_initiator=False)
599 def clear_with_event(
600 self, initiator: Optional[AttributeEventToken] = None
601 ) -> None:
602 """Empty the collection, firing a mutation event for each entity."""
604 if self.empty:
605 self._refuse_empty()
606 remover = self._data()._sa_remover
607 for item in list(self):
608 remover(item, _sa_initiator=initiator)
610 def clear_without_event(self) -> None:
611 """Empty the collection, firing no events."""
613 if self.empty:
614 self._refuse_empty()
615 remover = self._data()._sa_remover
616 for item in list(self):
617 remover(item, _sa_initiator=False)
619 def __iter__(self):
620 """Iterate over entities in the collection."""
622 return iter(self._data()._sa_iterator())
624 def __len__(self):
625 """Count entities in the collection."""
626 return len(list(self._data()._sa_iterator()))
628 def __bool__(self):
629 return True
631 def _fire_append_wo_mutation_event_bulk(
632 self, items, initiator=None, key=NO_KEY
633 ):
634 if not items:
635 return
637 if initiator is not False:
638 if self.invalidated:
639 self._warn_invalidated()
641 if self.empty:
642 self._reset_empty()
644 for item in items:
645 self.attr.fire_append_wo_mutation_event(
646 self.owner_state,
647 self.owner_state.dict,
648 item,
649 initiator,
650 key,
651 )
653 def fire_append_wo_mutation_event(self, item, initiator=None, key=NO_KEY):
654 """Notify that a entity is entering the collection but is already
655 present.
658 Initiator is a token owned by the InstrumentedAttribute that
659 initiated the membership mutation, and should be left as None
660 unless you are passing along an initiator value from a chained
661 operation.
663 .. versionadded:: 1.4.15
665 """
666 if initiator is not False:
667 if self.invalidated:
668 self._warn_invalidated()
670 if self.empty:
671 self._reset_empty()
673 return self.attr.fire_append_wo_mutation_event(
674 self.owner_state, self.owner_state.dict, item, initiator, key
675 )
676 else:
677 return item
679 def fire_append_event(self, item, initiator=None, key=NO_KEY):
680 """Notify that a entity has entered the collection.
682 Initiator is a token owned by the InstrumentedAttribute that
683 initiated the membership mutation, and should be left as None
684 unless you are passing along an initiator value from a chained
685 operation.
687 """
688 if initiator is not False:
689 if self.invalidated:
690 self._warn_invalidated()
692 if self.empty:
693 self._reset_empty()
695 return self.attr.fire_append_event(
696 self.owner_state, self.owner_state.dict, item, initiator, key
697 )
698 else:
699 return item
701 def _fire_remove_event_bulk(self, items, initiator=None, key=NO_KEY):
702 if not items:
703 return
705 if initiator is not False:
706 if self.invalidated:
707 self._warn_invalidated()
709 if self.empty:
710 self._reset_empty()
712 for item in items:
713 self.attr.fire_remove_event(
714 self.owner_state,
715 self.owner_state.dict,
716 item,
717 initiator,
718 key,
719 )
721 def fire_remove_event(self, item, initiator=None, key=NO_KEY):
722 """Notify that a entity has been removed from the collection.
724 Initiator is the InstrumentedAttribute that initiated the membership
725 mutation, and should be left as None unless you are passing along
726 an initiator value from a chained operation.
728 """
729 if initiator is not False:
730 if self.invalidated:
731 self._warn_invalidated()
733 if self.empty:
734 self._reset_empty()
736 self.attr.fire_remove_event(
737 self.owner_state, self.owner_state.dict, item, initiator, key
738 )
740 def fire_pre_remove_event(self, initiator=None, key=NO_KEY):
741 """Notify that an entity is about to be removed from the collection.
743 Only called if the entity cannot be removed after calling
744 fire_remove_event().
746 """
747 if self.invalidated:
748 self._warn_invalidated()
749 self.attr.fire_pre_remove_event(
750 self.owner_state,
751 self.owner_state.dict,
752 initiator=initiator,
753 key=key,
754 )
756 def __getstate__(self):
757 return {
758 "key": self._key,
759 "owner_state": self.owner_state,
760 "owner_cls": self.owner_state.class_,
761 "data": self.data,
762 "invalidated": self.invalidated,
763 "empty": self.empty,
764 }
766 def __setstate__(self, d):
767 self._key = d["key"]
768 self.owner_state = d["owner_state"]
770 # see note in constructor regarding this type: ignore
771 self._data = weakref.ref(d["data"]) # type: ignore
773 self._converter = d["data"]._sa_converter
774 d["data"]._sa_adapter = self
775 self.invalidated = d["invalidated"]
776 self.attr = getattr(d["owner_cls"], self._key).impl
777 self.empty = d.get("empty", False)
780def bulk_replace(values, existing_adapter, new_adapter, initiator=None):
781 """Load a new collection, firing events based on prior like membership.
783 Appends instances in ``values`` onto the ``new_adapter``. Events will be
784 fired for any instance not present in the ``existing_adapter``. Any
785 instances in ``existing_adapter`` not present in ``values`` will have
786 remove events fired upon them.
788 :param values: An iterable of collection member instances
790 :param existing_adapter: A :class:`.CollectionAdapter` of
791 instances to be replaced
793 :param new_adapter: An empty :class:`.CollectionAdapter`
794 to load with ``values``
797 """
799 assert isinstance(values, list)
801 idset = util.IdentitySet
802 existing_idset = idset(existing_adapter or ())
803 constants = existing_idset.intersection(values or ())
804 additions = idset(values or ()).difference(constants)
805 removals = existing_idset.difference(constants)
807 appender = new_adapter.bulk_appender()
809 for member in values or ():
810 if member in additions:
811 appender(member, _sa_initiator=initiator)
812 elif member in constants:
813 appender(member, _sa_initiator=False)
815 if existing_adapter:
816 existing_adapter._fire_append_wo_mutation_event_bulk(
817 constants, initiator=initiator
818 )
819 existing_adapter._fire_remove_event_bulk(removals, initiator=initiator)
822def prepare_instrumentation(
823 factory: Union[Type[Collection[Any]], _CollectionFactoryType],
824) -> _CollectionFactoryType:
825 """Prepare a callable for future use as a collection class factory.
827 Given a collection class factory (either a type or no-arg callable),
828 return another factory that will produce compatible instances when
829 called.
831 This function is responsible for converting collection_class=list
832 into the run-time behavior of collection_class=InstrumentedList.
834 """
836 impl_factory: _CollectionFactoryType
838 # Convert a builtin to 'Instrumented*'
839 if factory in __canned_instrumentation:
840 impl_factory = __canned_instrumentation[factory]
841 else:
842 impl_factory = cast(_CollectionFactoryType, factory)
844 cls: Union[_CollectionFactoryType, Type[Collection[Any]]]
846 # Create a specimen
847 cls = type(impl_factory())
849 # Did factory callable return a builtin?
850 if cls in __canned_instrumentation:
851 # if so, just convert.
852 # in previous major releases, this codepath wasn't working and was
853 # not covered by tests. prior to that it supplied a "wrapper"
854 # function that would return the class, though the rationale for this
855 # case is not known
856 impl_factory = __canned_instrumentation[cls]
857 cls = type(impl_factory())
859 # Instrument the class if needed.
860 if __instrumentation_mutex.acquire():
861 try:
862 if getattr(cls, "_sa_instrumented", None) != id(cls):
863 _instrument_class(cls)
864 finally:
865 __instrumentation_mutex.release()
867 return impl_factory
870def _instrument_class(cls):
871 """Modify methods in a class and install instrumentation."""
873 # In the normal call flow, a request for any of the 3 basic collection
874 # types is transformed into one of our trivial subclasses
875 # (e.g. InstrumentedList). Catch anything else that sneaks in here...
876 if cls.__module__ == "__builtin__":
877 raise sa_exc.ArgumentError(
878 "Can not instrument a built-in type. Use a "
879 "subclass, even a trivial one."
880 )
882 roles, methods = _locate_roles_and_methods(cls)
884 _setup_canned_roles(cls, roles, methods)
886 _assert_required_roles(cls, roles, methods)
888 _set_collection_attributes(cls, roles, methods)
891def _locate_roles_and_methods(cls):
892 """search for _sa_instrument_role-decorated methods in
893 method resolution order, assign to roles.
895 """
897 roles: Dict[str, str] = {}
898 methods: Dict[str, Tuple[Optional[str], Optional[int], Optional[str]]] = {}
900 for supercls in cls.__mro__:
901 for name, method in vars(supercls).items():
902 if not callable(method):
903 continue
905 # note role declarations
906 if hasattr(method, "_sa_instrument_role"):
907 role = method._sa_instrument_role
908 assert role in (
909 "appender",
910 "remover",
911 "iterator",
912 "converter",
913 )
914 roles.setdefault(role, name)
916 # transfer instrumentation requests from decorated function
917 # to the combined queue
918 before: Optional[Tuple[str, int]] = None
919 after: Optional[str] = None
921 if hasattr(method, "_sa_instrument_before"):
922 op, argument = method._sa_instrument_before
923 assert op in ("fire_append_event", "fire_remove_event")
924 before = op, argument
925 if hasattr(method, "_sa_instrument_after"):
926 op = method._sa_instrument_after
927 assert op in ("fire_append_event", "fire_remove_event")
928 after = op
929 if before:
930 methods[name] = before + (after,)
931 elif after:
932 methods[name] = None, None, after
933 return roles, methods
936def _setup_canned_roles(cls, roles, methods):
937 """see if this class has "canned" roles based on a known
938 collection type (dict, set, list). Apply those roles
939 as needed to the "roles" dictionary, and also
940 prepare "decorator" methods
942 """
943 collection_type = util.duck_type_collection(cls)
944 if collection_type in __interfaces:
945 assert collection_type is not None
946 canned_roles, decorators = __interfaces[collection_type]
947 for role, name in canned_roles.items():
948 roles.setdefault(role, name)
950 # apply ABC auto-decoration to methods that need it
951 for method, decorator in decorators.items():
952 fn = getattr(cls, method, None)
953 if (
954 fn
955 and method not in methods
956 and not hasattr(fn, "_sa_instrumented")
957 ):
958 setattr(cls, method, decorator(fn))
961def _assert_required_roles(cls, roles, methods):
962 """ensure all roles are present, and apply implicit instrumentation if
963 needed
965 """
966 if "appender" not in roles or not hasattr(cls, roles["appender"]):
967 raise sa_exc.ArgumentError(
968 "Type %s must elect an appender method to be "
969 "a collection class" % cls.__name__
970 )
971 elif roles["appender"] not in methods and not hasattr(
972 getattr(cls, roles["appender"]), "_sa_instrumented"
973 ):
974 methods[roles["appender"]] = ("fire_append_event", 1, None)
976 if "remover" not in roles or not hasattr(cls, roles["remover"]):
977 raise sa_exc.ArgumentError(
978 "Type %s must elect a remover method to be "
979 "a collection class" % cls.__name__
980 )
981 elif roles["remover"] not in methods and not hasattr(
982 getattr(cls, roles["remover"]), "_sa_instrumented"
983 ):
984 methods[roles["remover"]] = ("fire_remove_event", 1, None)
986 if "iterator" not in roles or not hasattr(cls, roles["iterator"]):
987 raise sa_exc.ArgumentError(
988 "Type %s must elect an iterator method to be "
989 "a collection class" % cls.__name__
990 )
993def _set_collection_attributes(cls, roles, methods):
994 """apply ad-hoc instrumentation from decorators, class-level defaults
995 and implicit role declarations
997 """
998 for method_name, (before, argument, after) in methods.items():
999 setattr(
1000 cls,
1001 method_name,
1002 _instrument_membership_mutator(
1003 getattr(cls, method_name), before, argument, after
1004 ),
1005 )
1006 # intern the role map
1007 for role, method_name in roles.items():
1008 setattr(cls, "_sa_%s" % role, getattr(cls, method_name))
1010 cls._sa_adapter = None
1012 if not hasattr(cls, "_sa_converter"):
1013 cls._sa_converter = None
1014 cls._sa_instrumented = id(cls)
1017def _instrument_membership_mutator(method, before, argument, after):
1018 """Route method args and/or return value through the collection
1019 adapter."""
1020 # This isn't smart enough to handle @adds(1) for 'def fn(self, (a, b))'
1021 if before:
1022 fn_args = list(
1023 util.flatten_iterator(inspect_getfullargspec(method)[0])
1024 )
1025 if isinstance(argument, int):
1026 pos_arg = argument
1027 named_arg = len(fn_args) > argument and fn_args[argument] or None
1028 else:
1029 if argument in fn_args:
1030 pos_arg = fn_args.index(argument)
1031 else:
1032 pos_arg = None
1033 named_arg = argument
1034 del fn_args
1036 def wrapper(*args, **kw):
1037 if before:
1038 if pos_arg is None:
1039 if named_arg not in kw:
1040 raise sa_exc.ArgumentError(
1041 "Missing argument %s" % argument
1042 )
1043 value = kw[named_arg]
1044 else:
1045 if len(args) > pos_arg:
1046 value = args[pos_arg]
1047 elif named_arg in kw:
1048 value = kw[named_arg]
1049 else:
1050 raise sa_exc.ArgumentError(
1051 "Missing argument %s" % argument
1052 )
1054 initiator = kw.pop("_sa_initiator", None)
1055 if initiator is False:
1056 executor = None
1057 else:
1058 executor = args[0]._sa_adapter
1060 if before and executor:
1061 getattr(executor, before)(value, initiator)
1063 if not after or not executor:
1064 return method(*args, **kw)
1065 else:
1066 res = method(*args, **kw)
1067 if res is not None:
1068 getattr(executor, after)(res, initiator)
1069 return res
1071 wrapper._sa_instrumented = True # type: ignore[attr-defined]
1072 if hasattr(method, "_sa_instrument_role"):
1073 wrapper._sa_instrument_role = method._sa_instrument_role # type: ignore[attr-defined] # noqa: E501
1074 wrapper.__name__ = method.__name__
1075 wrapper.__doc__ = method.__doc__
1076 return wrapper
1079def __set_wo_mutation(collection, item, _sa_initiator=None):
1080 """Run set wo mutation events.
1082 The collection is not mutated.
1084 """
1085 if _sa_initiator is not False:
1086 executor = collection._sa_adapter
1087 if executor:
1088 executor.fire_append_wo_mutation_event(
1089 item, _sa_initiator, key=None
1090 )
1093def __set(collection, item, _sa_initiator, key):
1094 """Run set events.
1096 This event always occurs before the collection is actually mutated.
1098 """
1100 if _sa_initiator is not False:
1101 executor = collection._sa_adapter
1102 if executor:
1103 item = executor.fire_append_event(item, _sa_initiator, key=key)
1104 return item
1107def __del(collection, item, _sa_initiator, key):
1108 """Run del events.
1110 This event occurs before the collection is actually mutated, *except*
1111 in the case of a pop operation, in which case it occurs afterwards.
1112 For pop operations, the __before_pop hook is called before the
1113 operation occurs.
1115 """
1116 if _sa_initiator is not False:
1117 executor = collection._sa_adapter
1118 if executor:
1119 executor.fire_remove_event(item, _sa_initiator, key=key)
1122def __before_pop(collection, _sa_initiator=None):
1123 """An event which occurs on a before a pop() operation occurs."""
1124 executor = collection._sa_adapter
1125 if executor:
1126 executor.fire_pre_remove_event(_sa_initiator)
1129def _list_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1130 """Tailored instrumentation wrappers for any list-like class."""
1132 def _tidy(fn):
1133 fn._sa_instrumented = True
1134 fn.__doc__ = getattr(list, fn.__name__).__doc__
1136 def append(fn):
1137 def append(self, item, _sa_initiator=None):
1138 item = __set(self, item, _sa_initiator, NO_KEY)
1139 fn(self, item)
1141 _tidy(append)
1142 return append
1144 def remove(fn):
1145 def remove(self, value, _sa_initiator=None):
1146 __del(self, value, _sa_initiator, NO_KEY)
1147 # testlib.pragma exempt:__eq__
1148 fn(self, value)
1150 _tidy(remove)
1151 return remove
1153 def insert(fn):
1154 def insert(self, index, value):
1155 value = __set(self, value, None, index)
1156 fn(self, index, value)
1158 _tidy(insert)
1159 return insert
1161 def __setitem__(fn):
1162 def __setitem__(self, index, value):
1163 if not isinstance(index, slice):
1164 existing = self[index]
1165 if existing is not None:
1166 __del(self, existing, None, index)
1167 value = __set(self, value, None, index)
1168 fn(self, index, value)
1169 else:
1170 # slice assignment requires __delitem__, insert, __len__
1171 step = index.step or 1
1172 start = index.start or 0
1173 if start < 0:
1174 start += len(self)
1175 if index.stop is not None:
1176 stop = index.stop
1177 else:
1178 stop = len(self)
1179 if stop < 0:
1180 stop += len(self)
1182 if step == 1:
1183 if value is self:
1184 return
1185 for i in range(start, stop, step):
1186 if len(self) > start:
1187 del self[start]
1189 for i, item in enumerate(value):
1190 self.insert(i + start, item)
1191 else:
1192 rng = list(range(start, stop, step))
1193 if len(value) != len(rng):
1194 raise ValueError(
1195 "attempt to assign sequence of size %s to "
1196 "extended slice of size %s"
1197 % (len(value), len(rng))
1198 )
1199 for i, item in zip(rng, value):
1200 self.__setitem__(i, item)
1202 _tidy(__setitem__)
1203 return __setitem__
1205 def __delitem__(fn):
1206 def __delitem__(self, index):
1207 if not isinstance(index, slice):
1208 item = self[index]
1209 __del(self, item, None, index)
1210 fn(self, index)
1211 else:
1212 # slice deletion requires __getslice__ and a slice-groking
1213 # __getitem__ for stepped deletion
1214 # note: not breaking this into atomic dels
1215 for item in self[index]:
1216 __del(self, item, None, index)
1217 fn(self, index)
1219 _tidy(__delitem__)
1220 return __delitem__
1222 def extend(fn):
1223 def extend(self, iterable):
1224 for value in list(iterable):
1225 self.append(value)
1227 _tidy(extend)
1228 return extend
1230 def __iadd__(fn):
1231 def __iadd__(self, iterable):
1232 # list.__iadd__ takes any iterable and seems to let TypeError
1233 # raise as-is instead of returning NotImplemented
1234 for value in list(iterable):
1235 self.append(value)
1236 return self
1238 _tidy(__iadd__)
1239 return __iadd__
1241 def pop(fn):
1242 def pop(self, index=-1):
1243 __before_pop(self)
1244 item = fn(self, index)
1245 __del(self, item, None, index)
1246 return item
1248 _tidy(pop)
1249 return pop
1251 def clear(fn):
1252 def clear(self, index=-1):
1253 for item in self:
1254 __del(self, item, None, index)
1255 fn(self)
1257 _tidy(clear)
1258 return clear
1260 # __imul__ : not wrapping this. all members of the collection are already
1261 # present, so no need to fire appends... wrapping it with an explicit
1262 # decorator is still possible, so events on *= can be had if they're
1263 # desired. hard to imagine a use case for __imul__, though.
1265 l = locals().copy()
1266 l.pop("_tidy")
1267 return l
1270def _dict_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1271 """Tailored instrumentation wrappers for any dict-like mapping class."""
1273 def _tidy(fn):
1274 fn._sa_instrumented = True
1275 fn.__doc__ = getattr(dict, fn.__name__).__doc__
1277 def __setitem__(fn):
1278 def __setitem__(self, key, value, _sa_initiator=None):
1279 if key in self:
1280 __del(self, self[key], _sa_initiator, key)
1281 value = __set(self, value, _sa_initiator, key)
1282 fn(self, key, value)
1284 _tidy(__setitem__)
1285 return __setitem__
1287 def __delitem__(fn):
1288 def __delitem__(self, key, _sa_initiator=None):
1289 if key in self:
1290 __del(self, self[key], _sa_initiator, key)
1291 fn(self, key)
1293 _tidy(__delitem__)
1294 return __delitem__
1296 def clear(fn):
1297 def clear(self):
1298 for key in self:
1299 __del(self, self[key], None, key)
1300 fn(self)
1302 _tidy(clear)
1303 return clear
1305 def pop(fn):
1306 def pop(self, key, default=NO_ARG):
1307 __before_pop(self)
1308 _to_del = key in self
1309 if default is NO_ARG:
1310 item = fn(self, key)
1311 else:
1312 item = fn(self, key, default)
1313 if _to_del:
1314 __del(self, item, None, key)
1315 return item
1317 _tidy(pop)
1318 return pop
1320 def popitem(fn):
1321 def popitem(self):
1322 __before_pop(self)
1323 item = fn(self)
1324 __del(self, item[1], None, 1)
1325 return item
1327 _tidy(popitem)
1328 return popitem
1330 def setdefault(fn):
1331 def setdefault(self, key, default=None):
1332 if key not in self:
1333 self.__setitem__(key, default)
1334 return default
1335 else:
1336 value = self.__getitem__(key)
1337 if value is default:
1338 __set_wo_mutation(self, value, None)
1340 return value
1342 _tidy(setdefault)
1343 return setdefault
1345 def update(fn):
1346 def update(self, __other=NO_ARG, **kw):
1347 if __other is not NO_ARG:
1348 if hasattr(__other, "keys"):
1349 for key in list(__other):
1350 if key not in self or self[key] is not __other[key]:
1351 self[key] = __other[key]
1352 else:
1353 __set_wo_mutation(self, __other[key], None)
1354 else:
1355 for key, value in __other:
1356 if key not in self or self[key] is not value:
1357 self[key] = value
1358 else:
1359 __set_wo_mutation(self, value, None)
1360 for key in kw:
1361 if key not in self or self[key] is not kw[key]:
1362 self[key] = kw[key]
1363 else:
1364 __set_wo_mutation(self, kw[key], None)
1366 _tidy(update)
1367 return update
1369 l = locals().copy()
1370 l.pop("_tidy")
1371 return l
1374_set_binop_bases = (set, frozenset)
1377def _set_binops_check_strict(self: Any, obj: Any) -> bool:
1378 """Allow only set, frozenset and self.__class__-derived
1379 objects in binops."""
1380 return isinstance(obj, _set_binop_bases + (self.__class__,))
1383def _set_binops_check_loose(self: Any, obj: Any) -> bool:
1384 """Allow anything set-like to participate in set binops."""
1385 return (
1386 isinstance(obj, _set_binop_bases + (self.__class__,))
1387 or util.duck_type_collection(obj) == set
1388 )
1391def _set_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1392 """Tailored instrumentation wrappers for any set-like class."""
1394 def _tidy(fn):
1395 fn._sa_instrumented = True
1396 fn.__doc__ = getattr(set, fn.__name__).__doc__
1398 def add(fn):
1399 def add(self, value, _sa_initiator=None):
1400 if value not in self:
1401 value = __set(self, value, _sa_initiator, NO_KEY)
1402 else:
1403 __set_wo_mutation(self, value, _sa_initiator)
1404 # testlib.pragma exempt:__hash__
1405 fn(self, value)
1407 _tidy(add)
1408 return add
1410 def discard(fn):
1411 def discard(self, value, _sa_initiator=None):
1412 # testlib.pragma exempt:__hash__
1413 if value in self:
1414 __del(self, value, _sa_initiator, NO_KEY)
1415 # testlib.pragma exempt:__hash__
1416 fn(self, value)
1418 _tidy(discard)
1419 return discard
1421 def remove(fn):
1422 def remove(self, value, _sa_initiator=None):
1423 # testlib.pragma exempt:__hash__
1424 if value in self:
1425 __del(self, value, _sa_initiator, NO_KEY)
1426 # testlib.pragma exempt:__hash__
1427 fn(self, value)
1429 _tidy(remove)
1430 return remove
1432 def pop(fn):
1433 def pop(self):
1434 __before_pop(self)
1435 item = fn(self)
1436 # for set in particular, we have no way to access the item
1437 # that will be popped before pop is called.
1438 __del(self, item, None, NO_KEY)
1439 return item
1441 _tidy(pop)
1442 return pop
1444 def clear(fn):
1445 def clear(self):
1446 for item in list(self):
1447 self.remove(item)
1449 _tidy(clear)
1450 return clear
1452 def update(fn):
1453 def update(self, value):
1454 for item in value:
1455 self.add(item)
1457 _tidy(update)
1458 return update
1460 def __ior__(fn):
1461 def __ior__(self, value):
1462 if not _set_binops_check_strict(self, value):
1463 return NotImplemented
1464 for item in value:
1465 self.add(item)
1466 return self
1468 _tidy(__ior__)
1469 return __ior__
1471 def difference_update(fn):
1472 def difference_update(self, value):
1473 for item in value:
1474 self.discard(item)
1476 _tidy(difference_update)
1477 return difference_update
1479 def __isub__(fn):
1480 def __isub__(self, value):
1481 if not _set_binops_check_strict(self, value):
1482 return NotImplemented
1483 for item in value:
1484 self.discard(item)
1485 return self
1487 _tidy(__isub__)
1488 return __isub__
1490 def intersection_update(fn):
1491 def intersection_update(self, other):
1492 want, have = self.intersection(other), set(self)
1493 remove, add = have - want, want - have
1495 for item in remove:
1496 self.remove(item)
1497 for item in add:
1498 self.add(item)
1500 _tidy(intersection_update)
1501 return intersection_update
1503 def __iand__(fn):
1504 def __iand__(self, other):
1505 if not _set_binops_check_strict(self, other):
1506 return NotImplemented
1507 want, have = self.intersection(other), set(self)
1508 remove, add = have - want, want - have
1510 for item in remove:
1511 self.remove(item)
1512 for item in add:
1513 self.add(item)
1514 return self
1516 _tidy(__iand__)
1517 return __iand__
1519 def symmetric_difference_update(fn):
1520 def symmetric_difference_update(self, other):
1521 want, have = self.symmetric_difference(other), set(self)
1522 remove, add = have - want, want - have
1524 for item in remove:
1525 self.remove(item)
1526 for item in add:
1527 self.add(item)
1529 _tidy(symmetric_difference_update)
1530 return symmetric_difference_update
1532 def __ixor__(fn):
1533 def __ixor__(self, other):
1534 if not _set_binops_check_strict(self, other):
1535 return NotImplemented
1536 want, have = self.symmetric_difference(other), set(self)
1537 remove, add = have - want, want - have
1539 for item in remove:
1540 self.remove(item)
1541 for item in add:
1542 self.add(item)
1543 return self
1545 _tidy(__ixor__)
1546 return __ixor__
1548 l = locals().copy()
1549 l.pop("_tidy")
1550 return l
1553class InstrumentedList(List[_T]):
1554 """An instrumented version of the built-in list."""
1557class InstrumentedSet(Set[_T]):
1558 """An instrumented version of the built-in set."""
1561class InstrumentedDict(Dict[_KT, _VT]):
1562 """An instrumented version of the built-in dict."""
1565__canned_instrumentation: util.immutabledict[Any, _CollectionFactoryType] = (
1566 util.immutabledict(
1567 {
1568 list: InstrumentedList,
1569 set: InstrumentedSet,
1570 dict: InstrumentedDict,
1571 }
1572 )
1573)
1575__interfaces: util.immutabledict[
1576 Any,
1577 Tuple[
1578 Dict[str, str],
1579 Dict[str, Callable[..., Any]],
1580 ],
1581] = util.immutabledict(
1582 {
1583 list: (
1584 {
1585 "appender": "append",
1586 "remover": "remove",
1587 "iterator": "__iter__",
1588 },
1589 _list_decorators(),
1590 ),
1591 set: (
1592 {"appender": "add", "remover": "remove", "iterator": "__iter__"},
1593 _set_decorators(),
1594 ),
1595 # decorators are required for dicts and object collections.
1596 dict: ({"iterator": "values"}, _dict_decorators()),
1597 }
1598)
1601def __go(lcls):
1602 global keyfunc_mapping, mapped_collection
1603 global column_keyed_dict, column_mapped_collection
1604 global MappedCollection, KeyFuncDict
1605 global attribute_keyed_dict, attribute_mapped_collection
1607 from .mapped_collection import keyfunc_mapping
1608 from .mapped_collection import column_keyed_dict
1609 from .mapped_collection import attribute_keyed_dict
1610 from .mapped_collection import KeyFuncDict
1612 from .mapped_collection import mapped_collection
1613 from .mapped_collection import column_mapped_collection
1614 from .mapped_collection import attribute_mapped_collection
1615 from .mapped_collection import MappedCollection
1617 # ensure instrumentation is associated with
1618 # these built-in classes; if a user-defined class
1619 # subclasses these and uses @internally_instrumented,
1620 # the superclass is otherwise not instrumented.
1621 # see [ticket:2406].
1622 _instrument_class(InstrumentedList)
1623 _instrument_class(InstrumentedSet)
1624 _instrument_class(KeyFuncDict)
1627__go(locals())