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 Protocol
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
137if typing.TYPE_CHECKING:
138 from .attributes import _CollectionAttributeImpl
139 from .attributes import AttributeEventToken
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]]
184class collection:
185 """Decorators for entity collection classes.
187 The decorators fall into two groups: annotations and interception recipes.
189 The annotating decorators (appender, remover, iterator,
190 internally_instrumented) indicate the method's purpose and take no
191 arguments. They are not written with parens::
193 @collection.appender
194 def append(self, append): ...
196 The recipe decorators all require parens, even those that take no
197 arguments::
199 @collection.adds("entity")
200 def insert(self, position, entity): ...
203 @collection.removes_return()
204 def popitem(self): ...
206 """
208 # Bundled as a class solely for ease of use: packaging, doc strings,
209 # importability.
211 @staticmethod
212 def appender(fn):
213 """Tag the method as the collection appender.
215 The appender method is called with one positional argument: the value
216 to append. The method will be automatically decorated with 'adds(1)'
217 if not already decorated::
219 @collection.appender
220 def add(self, append): ...
223 # or, equivalently
224 @collection.appender
225 @collection.adds(1)
226 def add(self, append): ...
229 # for mapping type, an 'append' may kick out a previous value
230 # that occupies that slot. consider d['a'] = 'foo'- any previous
231 # value in d['a'] is discarded.
232 @collection.appender
233 @collection.replaces(1)
234 def add(self, entity):
235 key = some_key_func(entity)
236 previous = None
237 if key in self:
238 previous = self[key]
239 self[key] = entity
240 return previous
242 If the value to append is not allowed in the collection, you may
243 raise an exception. Something to remember is that the appender
244 will be called for each object mapped by a database query. If the
245 database contains rows that violate your collection semantics, you
246 will need to get creative to fix the problem, as access via the
247 collection will not work.
249 If the appender method is internally instrumented, you must also
250 receive the keyword argument '_sa_initiator' and ensure its
251 promulgation to collection events.
253 """
254 fn._sa_instrument_role = "appender"
255 return fn
257 @staticmethod
258 def remover(fn):
259 """Tag the method as the collection remover.
261 The remover method is called with one positional argument: the value
262 to remove. The method will be automatically decorated with
263 :meth:`removes_return` if not already decorated::
265 @collection.remover
266 def zap(self, entity): ...
269 # or, equivalently
270 @collection.remover
271 @collection.removes_return()
272 def zap(self): ...
274 If the value to remove is not present in the collection, you may
275 raise an exception or return None to ignore the error.
277 If the remove method is internally instrumented, you must also
278 receive the keyword argument '_sa_initiator' and ensure its
279 promulgation to collection events.
281 """
282 fn._sa_instrument_role = "remover"
283 return fn
285 @staticmethod
286 def iterator(fn):
287 """Tag the method as the collection remover.
289 The iterator method is called with no arguments. It is expected to
290 return an iterator over all collection members::
292 @collection.iterator
293 def __iter__(self): ...
295 """
296 fn._sa_instrument_role = "iterator"
297 return fn
299 @staticmethod
300 def internally_instrumented(fn):
301 """Tag the method as instrumented.
303 This tag will prevent any decoration from being applied to the
304 method. Use this if you are orchestrating your own calls to
305 :func:`.collection_adapter` in one of the basic SQLAlchemy
306 interface methods, or to prevent an automatic ABC method
307 decoration from wrapping your implementation::
309 # normally an 'extend' method on a list-like class would be
310 # automatically intercepted and re-implemented in terms of
311 # SQLAlchemy events and append(). your implementation will
312 # never be called, unless:
313 @collection.internally_instrumented
314 def extend(self, items): ...
316 """
317 fn._sa_instrumented = True
318 return fn
320 @staticmethod
321 def adds(arg: int) -> Callable[[_FN], _FN]:
322 """Mark the method as adding an entity to the collection.
324 Adds "add to collection" handling to the method. The decorator
325 argument indicates which method argument holds the SQLAlchemy-relevant
326 value. Arguments can be specified positionally (i.e. integer) or by
327 name::
329 @collection.adds(1)
330 def push(self, item): ...
333 @collection.adds("entity")
334 def do_stuff(self, thing, entity=None): ...
336 """
338 def decorator(fn):
339 fn._sa_instrument_before = ("fire_append_event", arg)
340 return fn
342 return decorator
344 @staticmethod
345 def replaces(arg):
346 """Mark the method as replacing an entity in the collection.
348 Adds "add to collection" and "remove from collection" handling to
349 the method. The decorator argument indicates which method argument
350 holds the SQLAlchemy-relevant value to be added, and return value, if
351 any will be considered the value to remove.
353 Arguments can be specified positionally (i.e. integer) or by name::
355 @collection.replaces(2)
356 def __setitem__(self, index, item): ...
358 """
360 def decorator(fn):
361 fn._sa_instrument_before = ("fire_append_event", arg)
362 fn._sa_instrument_after = "fire_remove_event"
363 return fn
365 return decorator
367 @staticmethod
368 def removes(arg):
369 """Mark the method as removing an entity in the collection.
371 Adds "remove from collection" handling to the method. The decorator
372 argument indicates which method argument holds the SQLAlchemy-relevant
373 value to be removed. Arguments can be specified positionally (i.e.
374 integer) or by name::
376 @collection.removes(1)
377 def zap(self, item): ...
379 For methods where the value to remove is not known at call-time, use
380 collection.removes_return.
382 """
384 def decorator(fn):
385 fn._sa_instrument_before = ("fire_remove_event", arg)
386 return fn
388 return decorator
390 @staticmethod
391 def removes_return():
392 """Mark the method as removing an entity in the collection.
394 Adds "remove from collection" handling to the method. The return
395 value of the method, if any, is considered the value to remove. The
396 method arguments are not inspected::
398 @collection.removes_return()
399 def pop(self): ...
401 For methods where the value to remove is known at call-time, use
402 collection.remove.
404 """
406 def decorator(fn):
407 fn._sa_instrument_after = "fire_remove_event"
408 return fn
410 return decorator
413if TYPE_CHECKING:
415 def collection_adapter(collection: Collection[Any]) -> CollectionAdapter:
416 """Fetch the :class:`.CollectionAdapter` for a collection."""
418else:
419 collection_adapter = operator.attrgetter("_sa_adapter")
422class CollectionAdapter:
423 """Bridges between the ORM and arbitrary Python collections.
425 Proxies base-level collection operations (append, remove, iterate)
426 to the underlying Python collection, and emits add/remove events for
427 entities entering or leaving the collection.
429 The ORM uses :class:`.CollectionAdapter` exclusively for interaction with
430 entity collections.
433 """
435 __slots__ = (
436 "attr",
437 "_key",
438 "_data",
439 "owner_state",
440 "invalidated",
441 "empty",
442 )
444 attr: _CollectionAttributeImpl
445 _key: str
447 # this is actually a weakref; see note in constructor
448 _data: Callable[..., _AdaptedCollectionProtocol]
450 owner_state: InstanceState[Any]
451 invalidated: bool
452 empty: bool
454 def __init__(
455 self,
456 attr: _CollectionAttributeImpl,
457 owner_state: InstanceState[Any],
458 data: _AdaptedCollectionProtocol,
459 ):
460 self.attr = attr
461 self._key = attr.key
463 # this weakref stays referenced throughout the lifespan of
464 # CollectionAdapter. so while the weakref can return None, this
465 # is realistically only during garbage collection of this object, so
466 # we type this as a callable that returns _AdaptedCollectionProtocol
467 # in all cases.
468 self._data = weakref.ref(data) # type: ignore
470 self.owner_state = owner_state
471 data._sa_adapter = self
472 self.invalidated = False
473 self.empty = False
475 def _warn_invalidated(self) -> None:
476 util.warn("This collection has been invalidated.")
478 @property
479 def data(self) -> _AdaptedCollectionProtocol:
480 "The entity collection being adapted."
481 return self._data()
483 @property
484 def _referenced_by_owner(self) -> bool:
485 """return True if the owner state still refers to this collection.
487 This will return False within a bulk replace operation,
488 where this collection is the one being replaced.
490 """
491 return self.owner_state.dict[self._key] is self._data()
493 def bulk_appender(self):
494 return self._data()._sa_appender
496 def append_with_event(
497 self, item: Any, initiator: Optional[AttributeEventToken] = None
498 ) -> None:
499 """Add an entity to the collection, firing mutation events."""
501 self._data()._sa_appender(item, _sa_initiator=initiator)
503 def _set_empty(self, user_data):
504 assert (
505 not self.empty
506 ), "This collection adapter is already in the 'empty' state"
507 self.empty = True
508 self.owner_state._empty_collections[self._key] = user_data
510 def _reset_empty(self) -> None:
511 assert (
512 self.empty
513 ), "This collection adapter is not in the 'empty' state"
514 self.empty = False
515 self.owner_state.dict[self._key] = (
516 self.owner_state._empty_collections.pop(self._key)
517 )
519 def _refuse_empty(self) -> NoReturn:
520 raise sa_exc.InvalidRequestError(
521 "This is a special 'empty' collection which cannot accommodate "
522 "internal mutation operations"
523 )
525 def append_without_event(self, item: Any) -> None:
526 """Add or restore an entity to the collection, firing no events."""
528 if self.empty:
529 self._refuse_empty()
530 self._data()._sa_appender(item, _sa_initiator=False)
532 def append_multiple_without_event(self, items: Iterable[Any]) -> None:
533 """Add or restore an entity to the collection, firing no events."""
534 if self.empty:
535 self._refuse_empty()
536 appender = self._data()._sa_appender
537 for item in items:
538 appender(item, _sa_initiator=False)
540 def bulk_remover(self):
541 return self._data()._sa_remover
543 def remove_with_event(
544 self, item: Any, initiator: Optional[AttributeEventToken] = None
545 ) -> None:
546 """Remove an entity from the collection, firing mutation events."""
547 self._data()._sa_remover(item, _sa_initiator=initiator)
549 def remove_without_event(self, item: Any) -> None:
550 """Remove an entity from the collection, firing no events."""
551 if self.empty:
552 self._refuse_empty()
553 self._data()._sa_remover(item, _sa_initiator=False)
555 def clear_with_event(
556 self, initiator: Optional[AttributeEventToken] = None
557 ) -> None:
558 """Empty the collection, firing a mutation event for each entity."""
560 if self.empty:
561 self._refuse_empty()
562 remover = self._data()._sa_remover
563 for item in list(self):
564 remover(item, _sa_initiator=initiator)
566 def clear_without_event(self) -> None:
567 """Empty the collection, firing no events."""
569 if self.empty:
570 self._refuse_empty()
571 remover = self._data()._sa_remover
572 for item in list(self):
573 remover(item, _sa_initiator=False)
575 def __iter__(self):
576 """Iterate over entities in the collection."""
578 return iter(self._data()._sa_iterator())
580 def __len__(self):
581 """Count entities in the collection."""
582 return len(list(self._data()._sa_iterator()))
584 def __bool__(self):
585 return True
587 def _fire_append_wo_mutation_event_bulk(
588 self, items, initiator=None, key=NO_KEY
589 ):
590 if not items:
591 return
593 if initiator is not False:
594 if self.invalidated:
595 self._warn_invalidated()
597 if self.empty:
598 self._reset_empty()
600 for item in items:
601 self.attr.fire_append_wo_mutation_event(
602 self.owner_state,
603 self.owner_state.dict,
604 item,
605 initiator,
606 key,
607 )
609 def fire_append_wo_mutation_event(self, item, initiator=None, key=NO_KEY):
610 """Notify that a entity is entering the collection but is already
611 present.
614 Initiator is a token owned by the InstrumentedAttribute that
615 initiated the membership mutation, and should be left as None
616 unless you are passing along an initiator value from a chained
617 operation.
619 .. versionadded:: 1.4.15
621 """
622 if initiator is not False:
623 if self.invalidated:
624 self._warn_invalidated()
626 if self.empty:
627 self._reset_empty()
629 return self.attr.fire_append_wo_mutation_event(
630 self.owner_state, self.owner_state.dict, item, initiator, key
631 )
632 else:
633 return item
635 def fire_append_event(self, item, initiator=None, key=NO_KEY):
636 """Notify that a entity has entered the collection.
638 Initiator is a token owned by the InstrumentedAttribute that
639 initiated the membership mutation, and should be left as None
640 unless you are passing along an initiator value from a chained
641 operation.
643 """
644 if initiator is not False:
645 if self.invalidated:
646 self._warn_invalidated()
648 if self.empty:
649 self._reset_empty()
651 return self.attr.fire_append_event(
652 self.owner_state, self.owner_state.dict, item, initiator, key
653 )
654 else:
655 return item
657 def _fire_remove_event_bulk(self, items, initiator=None, key=NO_KEY):
658 if not items:
659 return
661 if initiator is not False:
662 if self.invalidated:
663 self._warn_invalidated()
665 if self.empty:
666 self._reset_empty()
668 for item in items:
669 self.attr.fire_remove_event(
670 self.owner_state,
671 self.owner_state.dict,
672 item,
673 initiator,
674 key,
675 )
677 def fire_remove_event(self, item, initiator=None, key=NO_KEY):
678 """Notify that a entity has been removed from the collection.
680 Initiator is the InstrumentedAttribute that initiated the membership
681 mutation, and should be left as None unless you are passing along
682 an initiator value from a chained operation.
684 """
685 if initiator is not False:
686 if self.invalidated:
687 self._warn_invalidated()
689 if self.empty:
690 self._reset_empty()
692 self.attr.fire_remove_event(
693 self.owner_state, self.owner_state.dict, item, initiator, key
694 )
696 def fire_pre_remove_event(self, initiator=None, key=NO_KEY):
697 """Notify that an entity is about to be removed from the collection.
699 Only called if the entity cannot be removed after calling
700 fire_remove_event().
702 """
703 if self.invalidated:
704 self._warn_invalidated()
705 self.attr.fire_pre_remove_event(
706 self.owner_state,
707 self.owner_state.dict,
708 initiator=initiator,
709 key=key,
710 )
712 def __getstate__(self):
713 return {
714 "key": self._key,
715 "owner_state": self.owner_state,
716 "owner_cls": self.owner_state.class_,
717 "data": self.data,
718 "invalidated": self.invalidated,
719 "empty": self.empty,
720 }
722 def __setstate__(self, d):
723 self._key = d["key"]
724 self.owner_state = d["owner_state"]
726 # see note in constructor regarding this type: ignore
727 self._data = weakref.ref(d["data"]) # type: ignore
729 d["data"]._sa_adapter = self
730 self.invalidated = d["invalidated"]
731 self.attr = getattr(d["owner_cls"], self._key).impl
732 self.empty = d.get("empty", False)
735def bulk_replace(values, existing_adapter, new_adapter, initiator=None):
736 """Load a new collection, firing events based on prior like membership.
738 Appends instances in ``values`` onto the ``new_adapter``. Events will be
739 fired for any instance not present in the ``existing_adapter``. Any
740 instances in ``existing_adapter`` not present in ``values`` will have
741 remove events fired upon them.
743 :param values: An iterable of collection member instances
745 :param existing_adapter: A :class:`.CollectionAdapter` of
746 instances to be replaced
748 :param new_adapter: An empty :class:`.CollectionAdapter`
749 to load with ``values``
752 """
754 assert isinstance(values, list)
756 idset = util.IdentitySet
757 existing_idset = idset(existing_adapter or ())
758 constants = existing_idset.intersection(values or ())
759 additions = idset(values or ()).difference(constants)
760 removals = existing_idset.difference(constants)
762 appender = new_adapter.bulk_appender()
764 for member in values or ():
765 if member in additions:
766 appender(member, _sa_initiator=initiator)
767 elif member in constants:
768 appender(member, _sa_initiator=False)
770 if existing_adapter:
771 existing_adapter._fire_append_wo_mutation_event_bulk(
772 constants, initiator=initiator
773 )
774 existing_adapter._fire_remove_event_bulk(removals, initiator=initiator)
777def _prepare_instrumentation(
778 factory: Union[Type[Collection[Any]], _CollectionFactoryType],
779) -> _CollectionFactoryType:
780 """Prepare a callable for future use as a collection class factory.
782 Given a collection class factory (either a type or no-arg callable),
783 return another factory that will produce compatible instances when
784 called.
786 This function is responsible for converting collection_class=list
787 into the run-time behavior of collection_class=InstrumentedList.
789 """
791 impl_factory: _CollectionFactoryType
793 # Convert a builtin to 'Instrumented*'
794 if factory in __canned_instrumentation:
795 impl_factory = __canned_instrumentation[factory]
796 else:
797 impl_factory = cast(_CollectionFactoryType, factory)
799 cls: Union[_CollectionFactoryType, Type[Collection[Any]]]
801 # Create a specimen
802 cls = type(impl_factory())
804 # Did factory callable return a builtin?
805 if cls in __canned_instrumentation:
806 # if so, just convert.
807 # in previous major releases, this codepath wasn't working and was
808 # not covered by tests. prior to that it supplied a "wrapper"
809 # function that would return the class, though the rationale for this
810 # case is not known
811 impl_factory = __canned_instrumentation[cls]
812 cls = type(impl_factory())
814 # Instrument the class if needed.
815 if __instrumentation_mutex.acquire():
816 try:
817 if getattr(cls, "_sa_instrumented", None) != id(cls):
818 _instrument_class(cls)
819 finally:
820 __instrumentation_mutex.release()
822 return impl_factory
825def _instrument_class(cls):
826 """Modify methods in a class and install instrumentation."""
828 # In the normal call flow, a request for any of the 3 basic collection
829 # types is transformed into one of our trivial subclasses
830 # (e.g. InstrumentedList). Catch anything else that sneaks in here...
831 if cls.__module__ == "__builtin__":
832 raise sa_exc.ArgumentError(
833 "Can not instrument a built-in type. Use a "
834 "subclass, even a trivial one."
835 )
837 roles, methods = _locate_roles_and_methods(cls)
839 _setup_canned_roles(cls, roles, methods)
841 _assert_required_roles(cls, roles, methods)
843 _set_collection_attributes(cls, roles, methods)
846def _locate_roles_and_methods(cls):
847 """search for _sa_instrument_role-decorated methods in
848 method resolution order, assign to roles.
850 """
852 roles: Dict[str, str] = {}
853 methods: Dict[str, Tuple[Optional[str], Optional[int], Optional[str]]] = {}
855 for supercls in cls.__mro__:
856 for name, method in vars(supercls).items():
857 if not callable(method):
858 continue
860 # note role declarations
861 if hasattr(method, "_sa_instrument_role"):
862 role = method._sa_instrument_role
863 assert role in ("appender", "remover", "iterator")
864 roles.setdefault(role, name)
866 # transfer instrumentation requests from decorated function
867 # to the combined queue
868 before: Optional[Tuple[str, int]] = None
869 after: Optional[str] = None
871 if hasattr(method, "_sa_instrument_before"):
872 op, argument = method._sa_instrument_before
873 assert op in ("fire_append_event", "fire_remove_event")
874 before = op, argument
875 if hasattr(method, "_sa_instrument_after"):
876 op = method._sa_instrument_after
877 assert op in ("fire_append_event", "fire_remove_event")
878 after = op
879 if before:
880 methods[name] = before + (after,)
881 elif after:
882 methods[name] = None, None, after
883 return roles, methods
886def _setup_canned_roles(cls, roles, methods):
887 """see if this class has "canned" roles based on a known
888 collection type (dict, set, list). Apply those roles
889 as needed to the "roles" dictionary, and also
890 prepare "decorator" methods
892 """
893 collection_type = util.duck_type_collection(cls)
894 if collection_type in __interfaces:
895 assert collection_type is not None
896 canned_roles, decorators = __interfaces[collection_type]
897 for role, name in canned_roles.items():
898 roles.setdefault(role, name)
900 # apply ABC auto-decoration to methods that need it
901 for method, decorator in decorators.items():
902 fn = getattr(cls, method, None)
903 if (
904 fn
905 and method not in methods
906 and not hasattr(fn, "_sa_instrumented")
907 ):
908 setattr(cls, method, decorator(fn))
911def _assert_required_roles(cls, roles, methods):
912 """ensure all roles are present, and apply implicit instrumentation if
913 needed
915 """
916 if "appender" not in roles or not hasattr(cls, roles["appender"]):
917 raise sa_exc.ArgumentError(
918 "Type %s must elect an appender method to be "
919 "a collection class" % cls.__name__
920 )
921 elif roles["appender"] not in methods and not hasattr(
922 getattr(cls, roles["appender"]), "_sa_instrumented"
923 ):
924 methods[roles["appender"]] = ("fire_append_event", 1, None)
926 if "remover" not in roles or not hasattr(cls, roles["remover"]):
927 raise sa_exc.ArgumentError(
928 "Type %s must elect a remover method to be "
929 "a collection class" % cls.__name__
930 )
931 elif roles["remover"] not in methods and not hasattr(
932 getattr(cls, roles["remover"]), "_sa_instrumented"
933 ):
934 methods[roles["remover"]] = ("fire_remove_event", 1, None)
936 if "iterator" not in roles or not hasattr(cls, roles["iterator"]):
937 raise sa_exc.ArgumentError(
938 "Type %s must elect an iterator method to be "
939 "a collection class" % cls.__name__
940 )
943def _set_collection_attributes(cls, roles, methods):
944 """apply ad-hoc instrumentation from decorators, class-level defaults
945 and implicit role declarations
947 """
948 for method_name, (before, argument, after) in methods.items():
949 setattr(
950 cls,
951 method_name,
952 _instrument_membership_mutator(
953 getattr(cls, method_name), before, argument, after
954 ),
955 )
956 # intern the role map
957 for role, method_name in roles.items():
958 setattr(cls, "_sa_%s" % role, getattr(cls, method_name))
960 cls._sa_adapter = None
962 cls._sa_instrumented = id(cls)
965def _instrument_membership_mutator(method, before, argument, after):
966 """Route method args and/or return value through the collection
967 adapter."""
968 # This isn't smart enough to handle @adds(1) for 'def fn(self, (a, b))'
969 if before:
970 fn_args = list(
971 util.flatten_iterator(inspect_getfullargspec(method)[0])
972 )
973 if isinstance(argument, int):
974 pos_arg = argument
975 named_arg = len(fn_args) > argument and fn_args[argument] or None
976 else:
977 if argument in fn_args:
978 pos_arg = fn_args.index(argument)
979 else:
980 pos_arg = None
981 named_arg = argument
982 del fn_args
984 def wrapper(*args, **kw):
985 if before:
986 if pos_arg is None:
987 if named_arg not in kw:
988 raise sa_exc.ArgumentError(
989 "Missing argument %s" % argument
990 )
991 value = kw[named_arg]
992 else:
993 if len(args) > pos_arg:
994 value = args[pos_arg]
995 elif named_arg in kw:
996 value = kw[named_arg]
997 else:
998 raise sa_exc.ArgumentError(
999 "Missing argument %s" % argument
1000 )
1002 initiator = kw.pop("_sa_initiator", None)
1003 if initiator is False:
1004 executor = None
1005 else:
1006 executor = args[0]._sa_adapter
1008 if before and executor:
1009 getattr(executor, before)(value, initiator)
1011 if not after or not executor:
1012 return method(*args, **kw)
1013 else:
1014 res = method(*args, **kw)
1015 if res is not None:
1016 getattr(executor, after)(res, initiator)
1017 return res
1019 wrapper._sa_instrumented = True # type: ignore[attr-defined]
1020 if hasattr(method, "_sa_instrument_role"):
1021 wrapper._sa_instrument_role = method._sa_instrument_role # type: ignore[attr-defined] # noqa: E501
1022 wrapper.__name__ = method.__name__
1023 wrapper.__doc__ = method.__doc__
1024 return wrapper
1027def __set_wo_mutation(collection, item, _sa_initiator=None):
1028 """Run set wo mutation events.
1030 The collection is not mutated.
1032 """
1033 if _sa_initiator is not False:
1034 executor = collection._sa_adapter
1035 if executor:
1036 executor.fire_append_wo_mutation_event(
1037 item, _sa_initiator, key=None
1038 )
1041def __set(collection, item, _sa_initiator, key):
1042 """Run set events.
1044 This event always occurs before the collection is actually mutated.
1046 """
1048 if _sa_initiator is not False:
1049 executor = collection._sa_adapter
1050 if executor:
1051 item = executor.fire_append_event(item, _sa_initiator, key=key)
1052 return item
1055def __del(collection, item, _sa_initiator, key):
1056 """Run del events.
1058 This event occurs before the collection is actually mutated, *except*
1059 in the case of a pop operation, in which case it occurs afterwards.
1060 For pop operations, the __before_pop hook is called before the
1061 operation occurs.
1063 """
1064 if _sa_initiator is not False:
1065 executor = collection._sa_adapter
1066 if executor:
1067 executor.fire_remove_event(item, _sa_initiator, key=key)
1070def __before_pop(collection, _sa_initiator=None):
1071 """An event which occurs on a before a pop() operation occurs."""
1072 executor = collection._sa_adapter
1073 if executor:
1074 executor.fire_pre_remove_event(_sa_initiator)
1077def _list_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1078 """Tailored instrumentation wrappers for any list-like class."""
1080 def _tidy(fn):
1081 fn._sa_instrumented = True
1082 fn.__doc__ = getattr(list, fn.__name__).__doc__
1084 def append(fn):
1085 def append(self, item, _sa_initiator=None):
1086 item = __set(self, item, _sa_initiator, NO_KEY)
1087 fn(self, item)
1089 _tidy(append)
1090 return append
1092 def remove(fn):
1093 def remove(self, value, _sa_initiator=None):
1094 __del(self, value, _sa_initiator, NO_KEY)
1095 # testlib.pragma exempt:__eq__
1096 fn(self, value)
1098 _tidy(remove)
1099 return remove
1101 def insert(fn):
1102 def insert(self, index, value):
1103 value = __set(self, value, None, index)
1104 fn(self, index, value)
1106 _tidy(insert)
1107 return insert
1109 def __setitem__(fn):
1110 def __setitem__(self, index, value):
1111 if not isinstance(index, slice):
1112 existing = self[index]
1113 if existing is not None:
1114 __del(self, existing, None, index)
1115 value = __set(self, value, None, index)
1116 fn(self, index, value)
1117 else:
1118 # slice assignment requires __delitem__, insert, __len__
1119 step = index.step or 1
1120 start = index.start or 0
1121 if start < 0:
1122 start += len(self)
1123 if index.stop is not None:
1124 stop = index.stop
1125 else:
1126 stop = len(self)
1127 if stop < 0:
1128 stop += len(self)
1130 if step == 1:
1131 if value is self:
1132 return
1133 for i in range(start, stop, step):
1134 if len(self) > start:
1135 del self[start]
1137 for i, item in enumerate(value):
1138 self.insert(i + start, item)
1139 else:
1140 rng = list(range(start, stop, step))
1141 if len(value) != len(rng):
1142 raise ValueError(
1143 "attempt to assign sequence of size %s to "
1144 "extended slice of size %s"
1145 % (len(value), len(rng))
1146 )
1147 for i, item in zip(rng, value):
1148 self.__setitem__(i, item)
1150 _tidy(__setitem__)
1151 return __setitem__
1153 def __delitem__(fn):
1154 def __delitem__(self, index):
1155 if not isinstance(index, slice):
1156 item = self[index]
1157 __del(self, item, None, index)
1158 fn(self, index)
1159 else:
1160 # slice deletion requires __getslice__ and a slice-groking
1161 # __getitem__ for stepped deletion
1162 # note: not breaking this into atomic dels
1163 for item in self[index]:
1164 __del(self, item, None, index)
1165 fn(self, index)
1167 _tidy(__delitem__)
1168 return __delitem__
1170 def extend(fn):
1171 def extend(self, iterable):
1172 for value in list(iterable):
1173 self.append(value)
1175 _tidy(extend)
1176 return extend
1178 def __iadd__(fn):
1179 def __iadd__(self, iterable):
1180 # list.__iadd__ takes any iterable and seems to let TypeError
1181 # raise as-is instead of returning NotImplemented
1182 for value in list(iterable):
1183 self.append(value)
1184 return self
1186 _tidy(__iadd__)
1187 return __iadd__
1189 def pop(fn):
1190 def pop(self, index=-1):
1191 __before_pop(self)
1192 item = fn(self, index)
1193 __del(self, item, None, index)
1194 return item
1196 _tidy(pop)
1197 return pop
1199 def clear(fn):
1200 def clear(self, index=-1):
1201 for item in self:
1202 __del(self, item, None, index)
1203 fn(self)
1205 _tidy(clear)
1206 return clear
1208 # __imul__ : not wrapping this. all members of the collection are already
1209 # present, so no need to fire appends... wrapping it with an explicit
1210 # decorator is still possible, so events on *= can be had if they're
1211 # desired. hard to imagine a use case for __imul__, though.
1213 l = locals().copy()
1214 l.pop("_tidy")
1215 return l
1218def _dict_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1219 """Tailored instrumentation wrappers for any dict-like mapping class."""
1221 def _tidy(fn):
1222 fn._sa_instrumented = True
1223 fn.__doc__ = getattr(dict, fn.__name__).__doc__
1225 def __setitem__(fn):
1226 def __setitem__(self, key, value, _sa_initiator=None):
1227 if key in self:
1228 __del(self, self[key], _sa_initiator, key)
1229 value = __set(self, value, _sa_initiator, key)
1230 fn(self, key, value)
1232 _tidy(__setitem__)
1233 return __setitem__
1235 def __delitem__(fn):
1236 def __delitem__(self, key, _sa_initiator=None):
1237 if key in self:
1238 __del(self, self[key], _sa_initiator, key)
1239 fn(self, key)
1241 _tidy(__delitem__)
1242 return __delitem__
1244 def clear(fn):
1245 def clear(self):
1246 for key in self:
1247 __del(self, self[key], None, key)
1248 fn(self)
1250 _tidy(clear)
1251 return clear
1253 def pop(fn):
1254 def pop(self, key, default=NO_ARG):
1255 __before_pop(self)
1256 _to_del = key in self
1257 if default is NO_ARG:
1258 item = fn(self, key)
1259 else:
1260 item = fn(self, key, default)
1261 if _to_del:
1262 __del(self, item, None, key)
1263 return item
1265 _tidy(pop)
1266 return pop
1268 def popitem(fn):
1269 def popitem(self):
1270 __before_pop(self)
1271 item = fn(self)
1272 __del(self, item[1], None, 1)
1273 return item
1275 _tidy(popitem)
1276 return popitem
1278 def setdefault(fn):
1279 def setdefault(self, key, default=None):
1280 if key not in self:
1281 self.__setitem__(key, default)
1282 return default
1283 else:
1284 value = self.__getitem__(key)
1285 if value is default:
1286 __set_wo_mutation(self, value, None)
1288 return value
1290 _tidy(setdefault)
1291 return setdefault
1293 def update(fn):
1294 def update(self, __other=NO_ARG, **kw):
1295 if __other is not NO_ARG:
1296 if hasattr(__other, "keys"):
1297 for key in list(__other):
1298 if key not in self or self[key] is not __other[key]:
1299 self[key] = __other[key]
1300 else:
1301 __set_wo_mutation(self, __other[key], None)
1302 else:
1303 for key, value in __other:
1304 if key not in self or self[key] is not value:
1305 self[key] = value
1306 else:
1307 __set_wo_mutation(self, value, None)
1308 for key in kw:
1309 if key not in self or self[key] is not kw[key]:
1310 self[key] = kw[key]
1311 else:
1312 __set_wo_mutation(self, kw[key], None)
1314 _tidy(update)
1315 return update
1317 l = locals().copy()
1318 l.pop("_tidy")
1319 return l
1322_set_binop_bases = (set, frozenset)
1325def _set_binops_check_strict(self: Any, obj: Any) -> bool:
1326 """Allow only set, frozenset and self.__class__-derived
1327 objects in binops."""
1328 return isinstance(obj, _set_binop_bases + (self.__class__,))
1331def _set_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1332 """Tailored instrumentation wrappers for any set-like class."""
1334 def _tidy(fn):
1335 fn._sa_instrumented = True
1336 fn.__doc__ = getattr(set, fn.__name__).__doc__
1338 def add(fn):
1339 def add(self, value, _sa_initiator=None):
1340 if value not in self:
1341 value = __set(self, value, _sa_initiator, NO_KEY)
1342 else:
1343 __set_wo_mutation(self, value, _sa_initiator)
1344 # testlib.pragma exempt:__hash__
1345 fn(self, value)
1347 _tidy(add)
1348 return add
1350 def discard(fn):
1351 def discard(self, value, _sa_initiator=None):
1352 # testlib.pragma exempt:__hash__
1353 if value in self:
1354 __del(self, value, _sa_initiator, NO_KEY)
1355 # testlib.pragma exempt:__hash__
1356 fn(self, value)
1358 _tidy(discard)
1359 return discard
1361 def remove(fn):
1362 def remove(self, value, _sa_initiator=None):
1363 # testlib.pragma exempt:__hash__
1364 if value in self:
1365 __del(self, value, _sa_initiator, NO_KEY)
1366 # testlib.pragma exempt:__hash__
1367 fn(self, value)
1369 _tidy(remove)
1370 return remove
1372 def pop(fn):
1373 def pop(self):
1374 __before_pop(self)
1375 item = fn(self)
1376 # for set in particular, we have no way to access the item
1377 # that will be popped before pop is called.
1378 __del(self, item, None, NO_KEY)
1379 return item
1381 _tidy(pop)
1382 return pop
1384 def clear(fn):
1385 def clear(self):
1386 for item in list(self):
1387 self.remove(item)
1389 _tidy(clear)
1390 return clear
1392 def update(fn):
1393 def update(self, value):
1394 for item in value:
1395 self.add(item)
1397 _tidy(update)
1398 return update
1400 def __ior__(fn):
1401 def __ior__(self, value):
1402 if not _set_binops_check_strict(self, value):
1403 return NotImplemented
1404 for item in value:
1405 self.add(item)
1406 return self
1408 _tidy(__ior__)
1409 return __ior__
1411 def difference_update(fn):
1412 def difference_update(self, value):
1413 for item in value:
1414 self.discard(item)
1416 _tidy(difference_update)
1417 return difference_update
1419 def __isub__(fn):
1420 def __isub__(self, value):
1421 if not _set_binops_check_strict(self, value):
1422 return NotImplemented
1423 for item in value:
1424 self.discard(item)
1425 return self
1427 _tidy(__isub__)
1428 return __isub__
1430 def intersection_update(fn):
1431 def intersection_update(self, other):
1432 want, have = self.intersection(other), set(self)
1433 remove, add = have - want, want - have
1435 for item in remove:
1436 self.remove(item)
1437 for item in add:
1438 self.add(item)
1440 _tidy(intersection_update)
1441 return intersection_update
1443 def __iand__(fn):
1444 def __iand__(self, other):
1445 if not _set_binops_check_strict(self, other):
1446 return NotImplemented
1447 want, have = self.intersection(other), set(self)
1448 remove, add = have - want, want - have
1450 for item in remove:
1451 self.remove(item)
1452 for item in add:
1453 self.add(item)
1454 return self
1456 _tidy(__iand__)
1457 return __iand__
1459 def symmetric_difference_update(fn):
1460 def symmetric_difference_update(self, other):
1461 want, have = self.symmetric_difference(other), set(self)
1462 remove, add = have - want, want - have
1464 for item in remove:
1465 self.remove(item)
1466 for item in add:
1467 self.add(item)
1469 _tidy(symmetric_difference_update)
1470 return symmetric_difference_update
1472 def __ixor__(fn):
1473 def __ixor__(self, other):
1474 if not _set_binops_check_strict(self, other):
1475 return NotImplemented
1476 want, have = self.symmetric_difference(other), set(self)
1477 remove, add = have - want, want - have
1479 for item in remove:
1480 self.remove(item)
1481 for item in add:
1482 self.add(item)
1483 return self
1485 _tidy(__ixor__)
1486 return __ixor__
1488 l = locals().copy()
1489 l.pop("_tidy")
1490 return l
1493class InstrumentedList(List[_T]):
1494 """An instrumented version of the built-in list."""
1497class InstrumentedSet(Set[_T]):
1498 """An instrumented version of the built-in set."""
1501class InstrumentedDict(Dict[_KT, _VT]):
1502 """An instrumented version of the built-in dict."""
1505__canned_instrumentation = cast(
1506 util.immutabledict[Any, _CollectionFactoryType],
1507 util.immutabledict(
1508 {
1509 list: InstrumentedList,
1510 set: InstrumentedSet,
1511 dict: InstrumentedDict,
1512 }
1513 ),
1514)
1516__interfaces: util.immutabledict[
1517 Any,
1518 Tuple[
1519 Dict[str, str],
1520 Dict[str, Callable[..., Any]],
1521 ],
1522] = util.immutabledict(
1523 {
1524 list: (
1525 {
1526 "appender": "append",
1527 "remover": "remove",
1528 "iterator": "__iter__",
1529 },
1530 _list_decorators(),
1531 ),
1532 set: (
1533 {"appender": "add", "remover": "remove", "iterator": "__iter__"},
1534 _set_decorators(),
1535 ),
1536 # decorators are required for dicts and object collections.
1537 dict: ({"iterator": "values"}, _dict_decorators()),
1538 }
1539)
1542def __go(lcls):
1543 global keyfunc_mapping, mapped_collection
1544 global column_keyed_dict, column_mapped_collection
1545 global MappedCollection, KeyFuncDict
1546 global attribute_keyed_dict, attribute_mapped_collection
1548 from .mapped_collection import keyfunc_mapping
1549 from .mapped_collection import column_keyed_dict
1550 from .mapped_collection import attribute_keyed_dict
1551 from .mapped_collection import KeyFuncDict
1553 from .mapped_collection import mapped_collection
1554 from .mapped_collection import column_mapped_collection
1555 from .mapped_collection import attribute_mapped_collection
1556 from .mapped_collection import MappedCollection
1558 # ensure instrumentation is associated with
1559 # these built-in classes; if a user-defined class
1560 # subclasses these and uses @internally_instrumented,
1561 # the superclass is otherwise not instrumented.
1562 # see [ticket:2406].
1563 _instrument_class(InstrumentedList)
1564 _instrument_class(InstrumentedSet)
1565 _instrument_class(KeyFuncDict)
1568__go(locals())