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-2024 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
24 class MyClass:
25 # ...
27 @collection.adds(1)
28 def store(self, item):
29 self.data.append(item)
31 @collection.removes_return()
32 def pop(self):
33 return self.data.pop()
36The second approach is a bundle of targeted decorators that wrap appropriate
37append and remove notifiers around the mutation methods present in the
38standard Python ``list``, ``set`` and ``dict`` interfaces. These could be
39specified in terms of generic decorator recipes, but are instead hand-tooled
40for increased efficiency. The targeted decorators occasionally implement
41adapter-like behavior, such as mapping bulk-set methods (``extend``,
42``update``, ``__setslice__``, etc.) into the series of atomic mutation events
43that the ORM requires.
45The targeted decorators are used internally for automatic instrumentation of
46entity collection classes. Every collection class goes through a
47transformation process roughly like so:
491. If the class is a built-in, substitute a trivial sub-class
502. Is this class already instrumented?
513. Add in generic decorators
524. Sniff out the collection interface through duck-typing
535. Add targeted decoration to any undecorated interface method
55This process modifies the class at runtime, decorating methods and adding some
56bookkeeping properties. This isn't possible (or desirable) for built-in
57classes like ``list``, so trivial sub-classes are substituted to hold
58decoration::
60 class InstrumentedList(list):
61 pass
63Collection classes can be specified in ``relationship(collection_class=)`` as
64types or a function that returns an instance. Collection classes are
65inspected and instrumented during the mapper compilation phase. The
66collection_class callable will be executed once to produce a specimen
67instance, and the type of that specimen will be instrumented. Functions that
68return built-in types like ``lists`` will be adapted to produce instrumented
69instances.
71When extending a known type like ``list``, additional decorations are not
72generally not needed. Odds are, the extension method will delegate to a
73method that's already instrumented. For example::
75 class QueueIsh(list):
76 def push(self, item):
77 self.append(item)
78 def shift(self):
79 return self.pop(0)
81There's no need to decorate these methods. ``append`` and ``pop`` are already
82instrumented as part of the ``list`` interface. Decorating them would fire
83duplicate events, which should be avoided.
85The targeted decoration tries not to rely on other methods in the underlying
86collection class, but some are unavoidable. Many depend on 'read' methods
87being present to properly instrument a 'write', for example, ``__setitem__``
88needs ``__getitem__``. "Bulk" methods like ``update`` and ``extend`` may also
89reimplemented in terms of atomic appends and removes, so the ``extend``
90decoration will actually perform many ``append`` operations and not call the
91underlying method at all.
93Tight control over bulk operation and the firing of events is also possible by
94implementing the instrumentation internally in your methods. The basic
95instrumentation package works under the general assumption that collection
96mutation will not raise unusual exceptions. If you want to closely
97orchestrate append and remove events with exception management, internal
98instrumentation may be the answer. Within your method,
99``collection_adapter(self)`` will retrieve an object that you can use for
100explicit control over triggering append and remove events.
102The owning object and :class:`.CollectionAttributeImpl` are also reachable
103through the adapter, allowing for some very sophisticated behavior.
105"""
106from __future__ import annotations
108import operator
109import threading
110import typing
111from typing import Any
112from typing import Callable
113from typing import cast
114from typing import Collection
115from typing import Dict
116from typing import Iterable
117from typing import List
118from typing import NoReturn
119from typing import Optional
120from typing import Protocol
121from typing import Set
122from typing import Tuple
123from typing import Type
124from typing import TYPE_CHECKING
125from typing import TypeVar
126from typing import Union
127import weakref
129from .base import NO_KEY
130from .. import exc as sa_exc
131from .. import util
132from ..sql.base import NO_ARG
133from ..util.compat import inspect_getfullargspec
135if typing.TYPE_CHECKING:
136 from .attributes import AttributeEventToken
137 from .attributes import CollectionAttributeImpl
138 from .mapped_collection import attribute_keyed_dict
139 from .mapped_collection import column_keyed_dict
140 from .mapped_collection import keyfunc_mapping
141 from .mapped_collection import KeyFuncDict # noqa: F401
142 from .state import InstanceState
145__all__ = [
146 "collection",
147 "collection_adapter",
148 "keyfunc_mapping",
149 "column_keyed_dict",
150 "attribute_keyed_dict",
151 "KeyFuncDict",
152 # old names in < 2.0
153 "mapped_collection",
154 "column_mapped_collection",
155 "attribute_mapped_collection",
156 "MappedCollection",
157]
159__instrumentation_mutex = threading.Lock()
162_CollectionFactoryType = Callable[[], "_AdaptedCollectionProtocol"]
164_T = TypeVar("_T", bound=Any)
165_KT = TypeVar("_KT", bound=Any)
166_VT = TypeVar("_VT", bound=Any)
167_COL = TypeVar("_COL", bound="Collection[Any]")
168_FN = TypeVar("_FN", bound="Callable[..., Any]")
171class _CollectionConverterProtocol(Protocol):
172 def __call__(self, collection: _COL) -> _COL: ...
175class _AdaptedCollectionProtocol(Protocol):
176 _sa_adapter: CollectionAdapter
177 _sa_appender: Callable[..., Any]
178 _sa_remover: Callable[..., Any]
179 _sa_iterator: Callable[..., Iterable[Any]]
180 _sa_converter: _CollectionConverterProtocol
183class collection:
184 """Decorators for entity collection classes.
186 The decorators fall into two groups: annotations and interception recipes.
188 The annotating decorators (appender, remover, iterator, converter,
189 internally_instrumented) indicate the method's purpose and take no
190 arguments. They are not written with parens::
192 @collection.appender
193 def append(self, append): ...
195 The recipe decorators all require parens, even those that take no
196 arguments::
198 @collection.adds('entity')
199 def insert(self, position, entity): ...
201 @collection.removes_return()
202 def popitem(self): ...
204 """
206 # Bundled as a class solely for ease of use: packaging, doc strings,
207 # importability.
209 @staticmethod
210 def appender(fn):
211 """Tag the method as the collection appender.
213 The appender method is called with one positional argument: the value
214 to append. The method will be automatically decorated with 'adds(1)'
215 if not already decorated::
217 @collection.appender
218 def add(self, append): ...
220 # or, equivalently
221 @collection.appender
222 @collection.adds(1)
223 def add(self, append): ...
225 # for mapping type, an 'append' may kick out a previous value
226 # that occupies that slot. consider d['a'] = 'foo'- any previous
227 # value in d['a'] is discarded.
228 @collection.appender
229 @collection.replaces(1)
230 def add(self, entity):
231 key = some_key_func(entity)
232 previous = None
233 if key in self:
234 previous = self[key]
235 self[key] = entity
236 return previous
238 If the value to append is not allowed in the collection, you may
239 raise an exception. Something to remember is that the appender
240 will be called for each object mapped by a database query. If the
241 database contains rows that violate your collection semantics, you
242 will need to get creative to fix the problem, as access via the
243 collection will not work.
245 If the appender method is internally instrumented, you must also
246 receive the keyword argument '_sa_initiator' and ensure its
247 promulgation to collection events.
249 """
250 fn._sa_instrument_role = "appender"
251 return fn
253 @staticmethod
254 def remover(fn):
255 """Tag the method as the collection remover.
257 The remover method is called with one positional argument: the value
258 to remove. The method will be automatically decorated with
259 :meth:`removes_return` if not already decorated::
261 @collection.remover
262 def zap(self, entity): ...
264 # or, equivalently
265 @collection.remover
266 @collection.removes_return()
267 def zap(self, ): ...
269 If the value to remove is not present in the collection, you may
270 raise an exception or return None to ignore the error.
272 If the remove method is internally instrumented, you must also
273 receive the keyword argument '_sa_initiator' and ensure its
274 promulgation to collection events.
276 """
277 fn._sa_instrument_role = "remover"
278 return fn
280 @staticmethod
281 def iterator(fn):
282 """Tag the method as the collection remover.
284 The iterator method is called with no arguments. It is expected to
285 return an iterator over all collection members::
287 @collection.iterator
288 def __iter__(self): ...
290 """
291 fn._sa_instrument_role = "iterator"
292 return fn
294 @staticmethod
295 def internally_instrumented(fn):
296 """Tag the method as instrumented.
298 This tag will prevent any decoration from being applied to the
299 method. Use this if you are orchestrating your own calls to
300 :func:`.collection_adapter` in one of the basic SQLAlchemy
301 interface methods, or to prevent an automatic ABC method
302 decoration from wrapping your implementation::
304 # normally an 'extend' method on a list-like class would be
305 # automatically intercepted and re-implemented in terms of
306 # SQLAlchemy events and append(). your implementation will
307 # never be called, unless:
308 @collection.internally_instrumented
309 def extend(self, items): ...
311 """
312 fn._sa_instrumented = True
313 return fn
315 @staticmethod
316 @util.deprecated(
317 "1.3",
318 "The :meth:`.collection.converter` handler is deprecated and will "
319 "be removed in a future release. Please refer to the "
320 ":class:`.AttributeEvents.bulk_replace` listener interface in "
321 "conjunction with the :func:`.event.listen` function.",
322 )
323 def converter(fn):
324 """Tag the method as the collection converter.
326 This optional method will be called when a collection is being
327 replaced entirely, as in::
329 myobj.acollection = [newvalue1, newvalue2]
331 The converter method will receive the object being assigned and should
332 return an iterable of values suitable for use by the ``appender``
333 method. A converter must not assign values or mutate the collection,
334 its sole job is to adapt the value the user provides into an iterable
335 of values for the ORM's use.
337 The default converter implementation will use duck-typing to do the
338 conversion. A dict-like collection will be convert into an iterable
339 of dictionary values, and other types will simply be iterated::
341 @collection.converter
342 def convert(self, other): ...
344 If the duck-typing of the object does not match the type of this
345 collection, a TypeError is raised.
347 Supply an implementation of this method if you want to expand the
348 range of possible types that can be assigned in bulk or perform
349 validation on the values about to be assigned.
351 """
352 fn._sa_instrument_role = "converter"
353 return fn
355 @staticmethod
356 def adds(arg):
357 """Mark the method as adding an entity to the collection.
359 Adds "add to collection" handling to the method. The decorator
360 argument indicates which method argument holds the SQLAlchemy-relevant
361 value. Arguments can be specified positionally (i.e. integer) or by
362 name::
364 @collection.adds(1)
365 def push(self, item): ...
367 @collection.adds('entity')
368 def do_stuff(self, thing, entity=None): ...
370 """
372 def decorator(fn):
373 fn._sa_instrument_before = ("fire_append_event", arg)
374 return fn
376 return decorator
378 @staticmethod
379 def replaces(arg):
380 """Mark the method as replacing an entity in the collection.
382 Adds "add to collection" and "remove from collection" handling to
383 the method. The decorator argument indicates which method argument
384 holds the SQLAlchemy-relevant value to be added, and return value, if
385 any will be considered the value to remove.
387 Arguments can be specified positionally (i.e. integer) or by name::
389 @collection.replaces(2)
390 def __setitem__(self, index, item): ...
392 """
394 def decorator(fn):
395 fn._sa_instrument_before = ("fire_append_event", arg)
396 fn._sa_instrument_after = "fire_remove_event"
397 return fn
399 return decorator
401 @staticmethod
402 def removes(arg):
403 """Mark the method as removing an entity in the collection.
405 Adds "remove from collection" handling to the method. The decorator
406 argument indicates which method argument holds the SQLAlchemy-relevant
407 value to be removed. Arguments can be specified positionally (i.e.
408 integer) or by name::
410 @collection.removes(1)
411 def zap(self, item): ...
413 For methods where the value to remove is not known at call-time, use
414 collection.removes_return.
416 """
418 def decorator(fn):
419 fn._sa_instrument_before = ("fire_remove_event", arg)
420 return fn
422 return decorator
424 @staticmethod
425 def removes_return():
426 """Mark the method as removing an entity in the collection.
428 Adds "remove from collection" handling to the method. The return
429 value of the method, if any, is considered the value to remove. The
430 method arguments are not inspected::
432 @collection.removes_return()
433 def pop(self): ...
435 For methods where the value to remove is known at call-time, use
436 collection.remove.
438 """
440 def decorator(fn):
441 fn._sa_instrument_after = "fire_remove_event"
442 return fn
444 return decorator
447if TYPE_CHECKING:
449 def collection_adapter(collection: Collection[Any]) -> CollectionAdapter:
450 """Fetch the :class:`.CollectionAdapter` for a collection."""
452else:
453 collection_adapter = operator.attrgetter("_sa_adapter")
456class CollectionAdapter:
457 """Bridges between the ORM and arbitrary Python collections.
459 Proxies base-level collection operations (append, remove, iterate)
460 to the underlying Python collection, and emits add/remove events for
461 entities entering or leaving the collection.
463 The ORM uses :class:`.CollectionAdapter` exclusively for interaction with
464 entity collections.
467 """
469 __slots__ = (
470 "attr",
471 "_key",
472 "_data",
473 "owner_state",
474 "_converter",
475 "invalidated",
476 "empty",
477 )
479 attr: CollectionAttributeImpl
480 _key: str
482 # this is actually a weakref; see note in constructor
483 _data: Callable[..., _AdaptedCollectionProtocol]
485 owner_state: InstanceState[Any]
486 _converter: _CollectionConverterProtocol
487 invalidated: bool
488 empty: bool
490 def __init__(
491 self,
492 attr: CollectionAttributeImpl,
493 owner_state: InstanceState[Any],
494 data: _AdaptedCollectionProtocol,
495 ):
496 self.attr = attr
497 self._key = attr.key
499 # this weakref stays referenced throughout the lifespan of
500 # CollectionAdapter. so while the weakref can return None, this
501 # is realistically only during garbage collection of this object, so
502 # we type this as a callable that returns _AdaptedCollectionProtocol
503 # in all cases.
504 self._data = weakref.ref(data) # type: ignore
506 self.owner_state = owner_state
507 data._sa_adapter = self
508 self._converter = data._sa_converter
509 self.invalidated = False
510 self.empty = False
512 def _warn_invalidated(self) -> None:
513 util.warn("This collection has been invalidated.")
515 @property
516 def data(self) -> _AdaptedCollectionProtocol:
517 "The entity collection being adapted."
518 return self._data()
520 @property
521 def _referenced_by_owner(self) -> bool:
522 """return True if the owner state still refers to this collection.
524 This will return False within a bulk replace operation,
525 where this collection is the one being replaced.
527 """
528 return self.owner_state.dict[self._key] is self._data()
530 def bulk_appender(self):
531 return self._data()._sa_appender
533 def append_with_event(
534 self, item: Any, initiator: Optional[AttributeEventToken] = None
535 ) -> None:
536 """Add an entity to the collection, firing mutation events."""
538 self._data()._sa_appender(item, _sa_initiator=initiator)
540 def _set_empty(self, user_data):
541 assert (
542 not self.empty
543 ), "This collection adapter is already in the 'empty' state"
544 self.empty = True
545 self.owner_state._empty_collections[self._key] = user_data
547 def _reset_empty(self) -> None:
548 assert (
549 self.empty
550 ), "This collection adapter is not in the 'empty' state"
551 self.empty = False
552 self.owner_state.dict[self._key] = (
553 self.owner_state._empty_collections.pop(self._key)
554 )
556 def _refuse_empty(self) -> NoReturn:
557 raise sa_exc.InvalidRequestError(
558 "This is a special 'empty' collection which cannot accommodate "
559 "internal mutation operations"
560 )
562 def append_without_event(self, item: Any) -> None:
563 """Add or restore an entity to the collection, firing no events."""
565 if self.empty:
566 self._refuse_empty()
567 self._data()._sa_appender(item, _sa_initiator=False)
569 def append_multiple_without_event(self, items: Iterable[Any]) -> None:
570 """Add or restore an entity to the collection, firing no events."""
571 if self.empty:
572 self._refuse_empty()
573 appender = self._data()._sa_appender
574 for item in items:
575 appender(item, _sa_initiator=False)
577 def bulk_remover(self):
578 return self._data()._sa_remover
580 def remove_with_event(
581 self, item: Any, initiator: Optional[AttributeEventToken] = None
582 ) -> None:
583 """Remove an entity from the collection, firing mutation events."""
584 self._data()._sa_remover(item, _sa_initiator=initiator)
586 def remove_without_event(self, item: Any) -> None:
587 """Remove an entity from the collection, firing no events."""
588 if self.empty:
589 self._refuse_empty()
590 self._data()._sa_remover(item, _sa_initiator=False)
592 def clear_with_event(
593 self, initiator: Optional[AttributeEventToken] = None
594 ) -> None:
595 """Empty the collection, firing a mutation event for each entity."""
597 if self.empty:
598 self._refuse_empty()
599 remover = self._data()._sa_remover
600 for item in list(self):
601 remover(item, _sa_initiator=initiator)
603 def clear_without_event(self) -> None:
604 """Empty the collection, firing no events."""
606 if self.empty:
607 self._refuse_empty()
608 remover = self._data()._sa_remover
609 for item in list(self):
610 remover(item, _sa_initiator=False)
612 def __iter__(self):
613 """Iterate over entities in the collection."""
615 return iter(self._data()._sa_iterator())
617 def __len__(self):
618 """Count entities in the collection."""
619 return len(list(self._data()._sa_iterator()))
621 def __bool__(self):
622 return True
624 def _fire_append_wo_mutation_event_bulk(
625 self, items, initiator=None, key=NO_KEY
626 ):
627 if not items:
628 return
630 if initiator is not False:
631 if self.invalidated:
632 self._warn_invalidated()
634 if self.empty:
635 self._reset_empty()
637 for item in items:
638 self.attr.fire_append_wo_mutation_event(
639 self.owner_state,
640 self.owner_state.dict,
641 item,
642 initiator,
643 key,
644 )
646 def fire_append_wo_mutation_event(self, item, initiator=None, key=NO_KEY):
647 """Notify that a entity is entering the collection but is already
648 present.
651 Initiator is a token owned by the InstrumentedAttribute that
652 initiated the membership mutation, and should be left as None
653 unless you are passing along an initiator value from a chained
654 operation.
656 .. versionadded:: 1.4.15
658 """
659 if initiator is not False:
660 if self.invalidated:
661 self._warn_invalidated()
663 if self.empty:
664 self._reset_empty()
666 return self.attr.fire_append_wo_mutation_event(
667 self.owner_state, self.owner_state.dict, item, initiator, key
668 )
669 else:
670 return item
672 def fire_append_event(self, item, initiator=None, key=NO_KEY):
673 """Notify that a entity has entered the collection.
675 Initiator is a token owned by the InstrumentedAttribute that
676 initiated the membership mutation, and should be left as None
677 unless you are passing along an initiator value from a chained
678 operation.
680 """
681 if initiator is not False:
682 if self.invalidated:
683 self._warn_invalidated()
685 if self.empty:
686 self._reset_empty()
688 return self.attr.fire_append_event(
689 self.owner_state, self.owner_state.dict, item, initiator, key
690 )
691 else:
692 return item
694 def _fire_remove_event_bulk(self, items, initiator=None, key=NO_KEY):
695 if not items:
696 return
698 if initiator is not False:
699 if self.invalidated:
700 self._warn_invalidated()
702 if self.empty:
703 self._reset_empty()
705 for item in items:
706 self.attr.fire_remove_event(
707 self.owner_state,
708 self.owner_state.dict,
709 item,
710 initiator,
711 key,
712 )
714 def fire_remove_event(self, item, initiator=None, key=NO_KEY):
715 """Notify that a entity has been removed from the collection.
717 Initiator is the InstrumentedAttribute that initiated the membership
718 mutation, and should be left as None unless you are passing along
719 an initiator value from a chained operation.
721 """
722 if initiator is not False:
723 if self.invalidated:
724 self._warn_invalidated()
726 if self.empty:
727 self._reset_empty()
729 self.attr.fire_remove_event(
730 self.owner_state, self.owner_state.dict, item, initiator, key
731 )
733 def fire_pre_remove_event(self, initiator=None, key=NO_KEY):
734 """Notify that an entity is about to be removed from the collection.
736 Only called if the entity cannot be removed after calling
737 fire_remove_event().
739 """
740 if self.invalidated:
741 self._warn_invalidated()
742 self.attr.fire_pre_remove_event(
743 self.owner_state,
744 self.owner_state.dict,
745 initiator=initiator,
746 key=key,
747 )
749 def __getstate__(self):
750 return {
751 "key": self._key,
752 "owner_state": self.owner_state,
753 "owner_cls": self.owner_state.class_,
754 "data": self.data,
755 "invalidated": self.invalidated,
756 "empty": self.empty,
757 }
759 def __setstate__(self, d):
760 self._key = d["key"]
761 self.owner_state = d["owner_state"]
763 # see note in constructor regarding this type: ignore
764 self._data = weakref.ref(d["data"]) # type: ignore
766 self._converter = d["data"]._sa_converter
767 d["data"]._sa_adapter = self
768 self.invalidated = d["invalidated"]
769 self.attr = getattr(d["owner_cls"], self._key).impl
770 self.empty = d.get("empty", False)
773def bulk_replace(values, existing_adapter, new_adapter, initiator=None):
774 """Load a new collection, firing events based on prior like membership.
776 Appends instances in ``values`` onto the ``new_adapter``. Events will be
777 fired for any instance not present in the ``existing_adapter``. Any
778 instances in ``existing_adapter`` not present in ``values`` will have
779 remove events fired upon them.
781 :param values: An iterable of collection member instances
783 :param existing_adapter: A :class:`.CollectionAdapter` of
784 instances to be replaced
786 :param new_adapter: An empty :class:`.CollectionAdapter`
787 to load with ``values``
790 """
792 assert isinstance(values, list)
794 idset = util.IdentitySet
795 existing_idset = idset(existing_adapter or ())
796 constants = existing_idset.intersection(values or ())
797 additions = idset(values or ()).difference(constants)
798 removals = existing_idset.difference(constants)
800 appender = new_adapter.bulk_appender()
802 for member in values or ():
803 if member in additions:
804 appender(member, _sa_initiator=initiator)
805 elif member in constants:
806 appender(member, _sa_initiator=False)
808 if existing_adapter:
809 existing_adapter._fire_append_wo_mutation_event_bulk(
810 constants, initiator=initiator
811 )
812 existing_adapter._fire_remove_event_bulk(removals, initiator=initiator)
815def prepare_instrumentation(
816 factory: Union[Type[Collection[Any]], _CollectionFactoryType],
817) -> _CollectionFactoryType:
818 """Prepare a callable for future use as a collection class factory.
820 Given a collection class factory (either a type or no-arg callable),
821 return another factory that will produce compatible instances when
822 called.
824 This function is responsible for converting collection_class=list
825 into the run-time behavior of collection_class=InstrumentedList.
827 """
829 impl_factory: _CollectionFactoryType
831 # Convert a builtin to 'Instrumented*'
832 if factory in __canned_instrumentation:
833 impl_factory = __canned_instrumentation[factory]
834 else:
835 impl_factory = cast(_CollectionFactoryType, factory)
837 cls: Union[_CollectionFactoryType, Type[Collection[Any]]]
839 # Create a specimen
840 cls = type(impl_factory())
842 # Did factory callable return a builtin?
843 if cls in __canned_instrumentation:
844 # if so, just convert.
845 # in previous major releases, this codepath wasn't working and was
846 # not covered by tests. prior to that it supplied a "wrapper"
847 # function that would return the class, though the rationale for this
848 # case is not known
849 impl_factory = __canned_instrumentation[cls]
850 cls = type(impl_factory())
852 # Instrument the class if needed.
853 if __instrumentation_mutex.acquire():
854 try:
855 if getattr(cls, "_sa_instrumented", None) != id(cls):
856 _instrument_class(cls)
857 finally:
858 __instrumentation_mutex.release()
860 return impl_factory
863def _instrument_class(cls):
864 """Modify methods in a class and install instrumentation."""
866 # In the normal call flow, a request for any of the 3 basic collection
867 # types is transformed into one of our trivial subclasses
868 # (e.g. InstrumentedList). Catch anything else that sneaks in here...
869 if cls.__module__ == "__builtin__":
870 raise sa_exc.ArgumentError(
871 "Can not instrument a built-in type. Use a "
872 "subclass, even a trivial one."
873 )
875 roles, methods = _locate_roles_and_methods(cls)
877 _setup_canned_roles(cls, roles, methods)
879 _assert_required_roles(cls, roles, methods)
881 _set_collection_attributes(cls, roles, methods)
884def _locate_roles_and_methods(cls):
885 """search for _sa_instrument_role-decorated methods in
886 method resolution order, assign to roles.
888 """
890 roles: Dict[str, str] = {}
891 methods: Dict[str, Tuple[Optional[str], Optional[int], Optional[str]]] = {}
893 for supercls in cls.__mro__:
894 for name, method in vars(supercls).items():
895 if not callable(method):
896 continue
898 # note role declarations
899 if hasattr(method, "_sa_instrument_role"):
900 role = method._sa_instrument_role
901 assert role in (
902 "appender",
903 "remover",
904 "iterator",
905 "converter",
906 )
907 roles.setdefault(role, name)
909 # transfer instrumentation requests from decorated function
910 # to the combined queue
911 before: Optional[Tuple[str, int]] = None
912 after: Optional[str] = None
914 if hasattr(method, "_sa_instrument_before"):
915 op, argument = method._sa_instrument_before
916 assert op in ("fire_append_event", "fire_remove_event")
917 before = op, argument
918 if hasattr(method, "_sa_instrument_after"):
919 op = method._sa_instrument_after
920 assert op in ("fire_append_event", "fire_remove_event")
921 after = op
922 if before:
923 methods[name] = before + (after,)
924 elif after:
925 methods[name] = None, None, after
926 return roles, methods
929def _setup_canned_roles(cls, roles, methods):
930 """see if this class has "canned" roles based on a known
931 collection type (dict, set, list). Apply those roles
932 as needed to the "roles" dictionary, and also
933 prepare "decorator" methods
935 """
936 collection_type = util.duck_type_collection(cls)
937 if collection_type in __interfaces:
938 assert collection_type is not None
939 canned_roles, decorators = __interfaces[collection_type]
940 for role, name in canned_roles.items():
941 roles.setdefault(role, name)
943 # apply ABC auto-decoration to methods that need it
944 for method, decorator in decorators.items():
945 fn = getattr(cls, method, None)
946 if (
947 fn
948 and method not in methods
949 and not hasattr(fn, "_sa_instrumented")
950 ):
951 setattr(cls, method, decorator(fn))
954def _assert_required_roles(cls, roles, methods):
955 """ensure all roles are present, and apply implicit instrumentation if
956 needed
958 """
959 if "appender" not in roles or not hasattr(cls, roles["appender"]):
960 raise sa_exc.ArgumentError(
961 "Type %s must elect an appender method to be "
962 "a collection class" % cls.__name__
963 )
964 elif roles["appender"] not in methods and not hasattr(
965 getattr(cls, roles["appender"]), "_sa_instrumented"
966 ):
967 methods[roles["appender"]] = ("fire_append_event", 1, None)
969 if "remover" not in roles or not hasattr(cls, roles["remover"]):
970 raise sa_exc.ArgumentError(
971 "Type %s must elect a remover method to be "
972 "a collection class" % cls.__name__
973 )
974 elif roles["remover"] not in methods and not hasattr(
975 getattr(cls, roles["remover"]), "_sa_instrumented"
976 ):
977 methods[roles["remover"]] = ("fire_remove_event", 1, None)
979 if "iterator" not in roles or not hasattr(cls, roles["iterator"]):
980 raise sa_exc.ArgumentError(
981 "Type %s must elect an iterator method to be "
982 "a collection class" % cls.__name__
983 )
986def _set_collection_attributes(cls, roles, methods):
987 """apply ad-hoc instrumentation from decorators, class-level defaults
988 and implicit role declarations
990 """
991 for method_name, (before, argument, after) in methods.items():
992 setattr(
993 cls,
994 method_name,
995 _instrument_membership_mutator(
996 getattr(cls, method_name), before, argument, after
997 ),
998 )
999 # intern the role map
1000 for role, method_name in roles.items():
1001 setattr(cls, "_sa_%s" % role, getattr(cls, method_name))
1003 cls._sa_adapter = None
1005 if not hasattr(cls, "_sa_converter"):
1006 cls._sa_converter = None
1007 cls._sa_instrumented = id(cls)
1010def _instrument_membership_mutator(method, before, argument, after):
1011 """Route method args and/or return value through the collection
1012 adapter."""
1013 # This isn't smart enough to handle @adds(1) for 'def fn(self, (a, b))'
1014 if before:
1015 fn_args = list(
1016 util.flatten_iterator(inspect_getfullargspec(method)[0])
1017 )
1018 if isinstance(argument, int):
1019 pos_arg = argument
1020 named_arg = len(fn_args) > argument and fn_args[argument] or None
1021 else:
1022 if argument in fn_args:
1023 pos_arg = fn_args.index(argument)
1024 else:
1025 pos_arg = None
1026 named_arg = argument
1027 del fn_args
1029 def wrapper(*args, **kw):
1030 if before:
1031 if pos_arg is None:
1032 if named_arg not in kw:
1033 raise sa_exc.ArgumentError(
1034 "Missing argument %s" % argument
1035 )
1036 value = kw[named_arg]
1037 else:
1038 if len(args) > pos_arg:
1039 value = args[pos_arg]
1040 elif named_arg in kw:
1041 value = kw[named_arg]
1042 else:
1043 raise sa_exc.ArgumentError(
1044 "Missing argument %s" % argument
1045 )
1047 initiator = kw.pop("_sa_initiator", None)
1048 if initiator is False:
1049 executor = None
1050 else:
1051 executor = args[0]._sa_adapter
1053 if before and executor:
1054 getattr(executor, before)(value, initiator)
1056 if not after or not executor:
1057 return method(*args, **kw)
1058 else:
1059 res = method(*args, **kw)
1060 if res is not None:
1061 getattr(executor, after)(res, initiator)
1062 return res
1064 wrapper._sa_instrumented = True # type: ignore[attr-defined]
1065 if hasattr(method, "_sa_instrument_role"):
1066 wrapper._sa_instrument_role = method._sa_instrument_role # type: ignore[attr-defined] # noqa: E501
1067 wrapper.__name__ = method.__name__
1068 wrapper.__doc__ = method.__doc__
1069 return wrapper
1072def __set_wo_mutation(collection, item, _sa_initiator=None):
1073 """Run set wo mutation events.
1075 The collection is not mutated.
1077 """
1078 if _sa_initiator is not False:
1079 executor = collection._sa_adapter
1080 if executor:
1081 executor.fire_append_wo_mutation_event(
1082 item, _sa_initiator, key=None
1083 )
1086def __set(collection, item, _sa_initiator, key):
1087 """Run set events.
1089 This event always occurs before the collection is actually mutated.
1091 """
1093 if _sa_initiator is not False:
1094 executor = collection._sa_adapter
1095 if executor:
1096 item = executor.fire_append_event(item, _sa_initiator, key=key)
1097 return item
1100def __del(collection, item, _sa_initiator, key):
1101 """Run del events.
1103 This event occurs before the collection is actually mutated, *except*
1104 in the case of a pop operation, in which case it occurs afterwards.
1105 For pop operations, the __before_pop hook is called before the
1106 operation occurs.
1108 """
1109 if _sa_initiator is not False:
1110 executor = collection._sa_adapter
1111 if executor:
1112 executor.fire_remove_event(item, _sa_initiator, key=key)
1115def __before_pop(collection, _sa_initiator=None):
1116 """An event which occurs on a before a pop() operation occurs."""
1117 executor = collection._sa_adapter
1118 if executor:
1119 executor.fire_pre_remove_event(_sa_initiator)
1122def _list_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1123 """Tailored instrumentation wrappers for any list-like class."""
1125 def _tidy(fn):
1126 fn._sa_instrumented = True
1127 fn.__doc__ = getattr(list, fn.__name__).__doc__
1129 def append(fn):
1130 def append(self, item, _sa_initiator=None):
1131 item = __set(self, item, _sa_initiator, NO_KEY)
1132 fn(self, item)
1134 _tidy(append)
1135 return append
1137 def remove(fn):
1138 def remove(self, value, _sa_initiator=None):
1139 __del(self, value, _sa_initiator, NO_KEY)
1140 # testlib.pragma exempt:__eq__
1141 fn(self, value)
1143 _tidy(remove)
1144 return remove
1146 def insert(fn):
1147 def insert(self, index, value):
1148 value = __set(self, value, None, index)
1149 fn(self, index, value)
1151 _tidy(insert)
1152 return insert
1154 def __setitem__(fn):
1155 def __setitem__(self, index, value):
1156 if not isinstance(index, slice):
1157 existing = self[index]
1158 if existing is not None:
1159 __del(self, existing, None, index)
1160 value = __set(self, value, None, index)
1161 fn(self, index, value)
1162 else:
1163 # slice assignment requires __delitem__, insert, __len__
1164 step = index.step or 1
1165 start = index.start or 0
1166 if start < 0:
1167 start += len(self)
1168 if index.stop is not None:
1169 stop = index.stop
1170 else:
1171 stop = len(self)
1172 if stop < 0:
1173 stop += len(self)
1175 if step == 1:
1176 if value is self:
1177 return
1178 for i in range(start, stop, step):
1179 if len(self) > start:
1180 del self[start]
1182 for i, item in enumerate(value):
1183 self.insert(i + start, item)
1184 else:
1185 rng = list(range(start, stop, step))
1186 if len(value) != len(rng):
1187 raise ValueError(
1188 "attempt to assign sequence of size %s to "
1189 "extended slice of size %s"
1190 % (len(value), len(rng))
1191 )
1192 for i, item in zip(rng, value):
1193 self.__setitem__(i, item)
1195 _tidy(__setitem__)
1196 return __setitem__
1198 def __delitem__(fn):
1199 def __delitem__(self, index):
1200 if not isinstance(index, slice):
1201 item = self[index]
1202 __del(self, item, None, index)
1203 fn(self, index)
1204 else:
1205 # slice deletion requires __getslice__ and a slice-groking
1206 # __getitem__ for stepped deletion
1207 # note: not breaking this into atomic dels
1208 for item in self[index]:
1209 __del(self, item, None, index)
1210 fn(self, index)
1212 _tidy(__delitem__)
1213 return __delitem__
1215 def extend(fn):
1216 def extend(self, iterable):
1217 for value in list(iterable):
1218 self.append(value)
1220 _tidy(extend)
1221 return extend
1223 def __iadd__(fn):
1224 def __iadd__(self, iterable):
1225 # list.__iadd__ takes any iterable and seems to let TypeError
1226 # raise as-is instead of returning NotImplemented
1227 for value in list(iterable):
1228 self.append(value)
1229 return self
1231 _tidy(__iadd__)
1232 return __iadd__
1234 def pop(fn):
1235 def pop(self, index=-1):
1236 __before_pop(self)
1237 item = fn(self, index)
1238 __del(self, item, None, index)
1239 return item
1241 _tidy(pop)
1242 return pop
1244 def clear(fn):
1245 def clear(self, index=-1):
1246 for item in self:
1247 __del(self, item, None, index)
1248 fn(self)
1250 _tidy(clear)
1251 return clear
1253 # __imul__ : not wrapping this. all members of the collection are already
1254 # present, so no need to fire appends... wrapping it with an explicit
1255 # decorator is still possible, so events on *= can be had if they're
1256 # desired. hard to imagine a use case for __imul__, though.
1258 l = locals().copy()
1259 l.pop("_tidy")
1260 return l
1263def _dict_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1264 """Tailored instrumentation wrappers for any dict-like mapping class."""
1266 def _tidy(fn):
1267 fn._sa_instrumented = True
1268 fn.__doc__ = getattr(dict, fn.__name__).__doc__
1270 def __setitem__(fn):
1271 def __setitem__(self, key, value, _sa_initiator=None):
1272 if key in self:
1273 __del(self, self[key], _sa_initiator, key)
1274 value = __set(self, value, _sa_initiator, key)
1275 fn(self, key, value)
1277 _tidy(__setitem__)
1278 return __setitem__
1280 def __delitem__(fn):
1281 def __delitem__(self, key, _sa_initiator=None):
1282 if key in self:
1283 __del(self, self[key], _sa_initiator, key)
1284 fn(self, key)
1286 _tidy(__delitem__)
1287 return __delitem__
1289 def clear(fn):
1290 def clear(self):
1291 for key in self:
1292 __del(self, self[key], None, key)
1293 fn(self)
1295 _tidy(clear)
1296 return clear
1298 def pop(fn):
1299 def pop(self, key, default=NO_ARG):
1300 __before_pop(self)
1301 _to_del = key in self
1302 if default is NO_ARG:
1303 item = fn(self, key)
1304 else:
1305 item = fn(self, key, default)
1306 if _to_del:
1307 __del(self, item, None, key)
1308 return item
1310 _tidy(pop)
1311 return pop
1313 def popitem(fn):
1314 def popitem(self):
1315 __before_pop(self)
1316 item = fn(self)
1317 __del(self, item[1], None, 1)
1318 return item
1320 _tidy(popitem)
1321 return popitem
1323 def setdefault(fn):
1324 def setdefault(self, key, default=None):
1325 if key not in self:
1326 self.__setitem__(key, default)
1327 return default
1328 else:
1329 value = self.__getitem__(key)
1330 if value is default:
1331 __set_wo_mutation(self, value, None)
1333 return value
1335 _tidy(setdefault)
1336 return setdefault
1338 def update(fn):
1339 def update(self, __other=NO_ARG, **kw):
1340 if __other is not NO_ARG:
1341 if hasattr(__other, "keys"):
1342 for key in list(__other):
1343 if key not in self or self[key] is not __other[key]:
1344 self[key] = __other[key]
1345 else:
1346 __set_wo_mutation(self, __other[key], None)
1347 else:
1348 for key, value in __other:
1349 if key not in self or self[key] is not value:
1350 self[key] = value
1351 else:
1352 __set_wo_mutation(self, value, None)
1353 for key in kw:
1354 if key not in self or self[key] is not kw[key]:
1355 self[key] = kw[key]
1356 else:
1357 __set_wo_mutation(self, kw[key], None)
1359 _tidy(update)
1360 return update
1362 l = locals().copy()
1363 l.pop("_tidy")
1364 return l
1367_set_binop_bases = (set, frozenset)
1370def _set_binops_check_strict(self: Any, obj: Any) -> bool:
1371 """Allow only set, frozenset and self.__class__-derived
1372 objects in binops."""
1373 return isinstance(obj, _set_binop_bases + (self.__class__,))
1376def _set_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1377 """Tailored instrumentation wrappers for any set-like class."""
1379 def _tidy(fn):
1380 fn._sa_instrumented = True
1381 fn.__doc__ = getattr(set, fn.__name__).__doc__
1383 def add(fn):
1384 def add(self, value, _sa_initiator=None):
1385 if value not in self:
1386 value = __set(self, value, _sa_initiator, NO_KEY)
1387 else:
1388 __set_wo_mutation(self, value, _sa_initiator)
1389 # testlib.pragma exempt:__hash__
1390 fn(self, value)
1392 _tidy(add)
1393 return add
1395 def discard(fn):
1396 def discard(self, value, _sa_initiator=None):
1397 # testlib.pragma exempt:__hash__
1398 if value in self:
1399 __del(self, value, _sa_initiator, NO_KEY)
1400 # testlib.pragma exempt:__hash__
1401 fn(self, value)
1403 _tidy(discard)
1404 return discard
1406 def remove(fn):
1407 def remove(self, value, _sa_initiator=None):
1408 # testlib.pragma exempt:__hash__
1409 if value in self:
1410 __del(self, value, _sa_initiator, NO_KEY)
1411 # testlib.pragma exempt:__hash__
1412 fn(self, value)
1414 _tidy(remove)
1415 return remove
1417 def pop(fn):
1418 def pop(self):
1419 __before_pop(self)
1420 item = fn(self)
1421 # for set in particular, we have no way to access the item
1422 # that will be popped before pop is called.
1423 __del(self, item, None, NO_KEY)
1424 return item
1426 _tidy(pop)
1427 return pop
1429 def clear(fn):
1430 def clear(self):
1431 for item in list(self):
1432 self.remove(item)
1434 _tidy(clear)
1435 return clear
1437 def update(fn):
1438 def update(self, value):
1439 for item in value:
1440 self.add(item)
1442 _tidy(update)
1443 return update
1445 def __ior__(fn):
1446 def __ior__(self, value):
1447 if not _set_binops_check_strict(self, value):
1448 return NotImplemented
1449 for item in value:
1450 self.add(item)
1451 return self
1453 _tidy(__ior__)
1454 return __ior__
1456 def difference_update(fn):
1457 def difference_update(self, value):
1458 for item in value:
1459 self.discard(item)
1461 _tidy(difference_update)
1462 return difference_update
1464 def __isub__(fn):
1465 def __isub__(self, value):
1466 if not _set_binops_check_strict(self, value):
1467 return NotImplemented
1468 for item in value:
1469 self.discard(item)
1470 return self
1472 _tidy(__isub__)
1473 return __isub__
1475 def intersection_update(fn):
1476 def intersection_update(self, other):
1477 want, have = self.intersection(other), set(self)
1478 remove, add = have - want, want - have
1480 for item in remove:
1481 self.remove(item)
1482 for item in add:
1483 self.add(item)
1485 _tidy(intersection_update)
1486 return intersection_update
1488 def __iand__(fn):
1489 def __iand__(self, other):
1490 if not _set_binops_check_strict(self, other):
1491 return NotImplemented
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)
1499 return self
1501 _tidy(__iand__)
1502 return __iand__
1504 def symmetric_difference_update(fn):
1505 def symmetric_difference_update(self, other):
1506 want, have = self.symmetric_difference(other), set(self)
1507 remove, add = have - want, want - have
1509 for item in remove:
1510 self.remove(item)
1511 for item in add:
1512 self.add(item)
1514 _tidy(symmetric_difference_update)
1515 return symmetric_difference_update
1517 def __ixor__(fn):
1518 def __ixor__(self, other):
1519 if not _set_binops_check_strict(self, other):
1520 return NotImplemented
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)
1528 return self
1530 _tidy(__ixor__)
1531 return __ixor__
1533 l = locals().copy()
1534 l.pop("_tidy")
1535 return l
1538class InstrumentedList(List[_T]):
1539 """An instrumented version of the built-in list."""
1542class InstrumentedSet(Set[_T]):
1543 """An instrumented version of the built-in set."""
1546class InstrumentedDict(Dict[_KT, _VT]):
1547 """An instrumented version of the built-in dict."""
1550__canned_instrumentation = cast(
1551 util.immutabledict[Any, _CollectionFactoryType],
1552 util.immutabledict(
1553 {
1554 list: InstrumentedList,
1555 set: InstrumentedSet,
1556 dict: InstrumentedDict,
1557 }
1558 ),
1559)
1561__interfaces: util.immutabledict[
1562 Any,
1563 Tuple[
1564 Dict[str, str],
1565 Dict[str, Callable[..., Any]],
1566 ],
1567] = util.immutabledict(
1568 {
1569 list: (
1570 {
1571 "appender": "append",
1572 "remover": "remove",
1573 "iterator": "__iter__",
1574 },
1575 _list_decorators(),
1576 ),
1577 set: (
1578 {"appender": "add", "remover": "remove", "iterator": "__iter__"},
1579 _set_decorators(),
1580 ),
1581 # decorators are required for dicts and object collections.
1582 dict: ({"iterator": "values"}, _dict_decorators()),
1583 }
1584)
1587def __go(lcls):
1588 global keyfunc_mapping, mapped_collection
1589 global column_keyed_dict, column_mapped_collection
1590 global MappedCollection, KeyFuncDict
1591 global attribute_keyed_dict, attribute_mapped_collection
1593 from .mapped_collection import keyfunc_mapping
1594 from .mapped_collection import column_keyed_dict
1595 from .mapped_collection import attribute_keyed_dict
1596 from .mapped_collection import KeyFuncDict
1598 from .mapped_collection import mapped_collection
1599 from .mapped_collection import column_mapped_collection
1600 from .mapped_collection import attribute_mapped_collection
1601 from .mapped_collection import MappedCollection
1603 # ensure instrumentation is associated with
1604 # these built-in classes; if a user-defined class
1605 # subclasses these and uses @internally_instrumented,
1606 # the superclass is otherwise not instrumented.
1607 # see [ticket:2406].
1608 _instrument_class(InstrumentedList)
1609 _instrument_class(InstrumentedSet)
1610 _instrument_class(KeyFuncDict)
1613__go(locals())