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 "column_keyed_dict",
152 "attribute_keyed_dict",
153 "MappedCollection",
154 "KeyFuncDict",
155]
157__instrumentation_mutex = threading.Lock()
160_CollectionFactoryType = Callable[[], "_AdaptedCollectionProtocol"]
162_T = TypeVar("_T", bound=Any)
163_KT = TypeVar("_KT", bound=Any)
164_VT = TypeVar("_VT", bound=Any)
165_COL = TypeVar("_COL", bound="Collection[Any]")
166_FN = TypeVar("_FN", bound="Callable[..., Any]")
169class _CollectionConverterProtocol(Protocol):
170 def __call__(self, collection: _COL) -> _COL: ...
173class _AdaptedCollectionProtocol(Protocol):
174 _sa_adapter: CollectionAdapter
175 _sa_appender: Callable[..., Any]
176 _sa_remover: Callable[..., Any]
177 _sa_iterator: Callable[..., Iterable[Any]]
178 _sa_converter: _CollectionConverterProtocol
181class collection:
182 """Decorators for entity collection classes.
184 The decorators fall into two groups: annotations and interception recipes.
186 The annotating decorators (appender, remover, iterator, converter,
187 internally_instrumented) indicate the method's purpose and take no
188 arguments. They are not written with parens::
190 @collection.appender
191 def append(self, append): ...
193 The recipe decorators all require parens, even those that take no
194 arguments::
196 @collection.adds('entity')
197 def insert(self, position, entity): ...
199 @collection.removes_return()
200 def popitem(self): ...
202 """
204 # Bundled as a class solely for ease of use: packaging, doc strings,
205 # importability.
207 @staticmethod
208 def appender(fn):
209 """Tag the method as the collection appender.
211 The appender method is called with one positional argument: the value
212 to append. The method will be automatically decorated with 'adds(1)'
213 if not already decorated::
215 @collection.appender
216 def add(self, append): ...
218 # or, equivalently
219 @collection.appender
220 @collection.adds(1)
221 def add(self, append): ...
223 # for mapping type, an 'append' may kick out a previous value
224 # that occupies that slot. consider d['a'] = 'foo'- any previous
225 # value in d['a'] is discarded.
226 @collection.appender
227 @collection.replaces(1)
228 def add(self, entity):
229 key = some_key_func(entity)
230 previous = None
231 if key in self:
232 previous = self[key]
233 self[key] = entity
234 return previous
236 If the value to append is not allowed in the collection, you may
237 raise an exception. Something to remember is that the appender
238 will be called for each object mapped by a database query. If the
239 database contains rows that violate your collection semantics, you
240 will need to get creative to fix the problem, as access via the
241 collection will not work.
243 If the appender method is internally instrumented, you must also
244 receive the keyword argument '_sa_initiator' and ensure its
245 promulgation to collection events.
247 """
248 fn._sa_instrument_role = "appender"
249 return fn
251 @staticmethod
252 def remover(fn):
253 """Tag the method as the collection remover.
255 The remover method is called with one positional argument: the value
256 to remove. The method will be automatically decorated with
257 :meth:`removes_return` if not already decorated::
259 @collection.remover
260 def zap(self, entity): ...
262 # or, equivalently
263 @collection.remover
264 @collection.removes_return()
265 def zap(self, ): ...
267 If the value to remove is not present in the collection, you may
268 raise an exception or return None to ignore the error.
270 If the remove method is internally instrumented, you must also
271 receive the keyword argument '_sa_initiator' and ensure its
272 promulgation to collection events.
274 """
275 fn._sa_instrument_role = "remover"
276 return fn
278 @staticmethod
279 def iterator(fn):
280 """Tag the method as the collection remover.
282 The iterator method is called with no arguments. It is expected to
283 return an iterator over all collection members::
285 @collection.iterator
286 def __iter__(self): ...
288 """
289 fn._sa_instrument_role = "iterator"
290 return fn
292 @staticmethod
293 def internally_instrumented(fn):
294 """Tag the method as instrumented.
296 This tag will prevent any decoration from being applied to the
297 method. Use this if you are orchestrating your own calls to
298 :func:`.collection_adapter` in one of the basic SQLAlchemy
299 interface methods, or to prevent an automatic ABC method
300 decoration from wrapping your implementation::
302 # normally an 'extend' method on a list-like class would be
303 # automatically intercepted and re-implemented in terms of
304 # SQLAlchemy events and append(). your implementation will
305 # never be called, unless:
306 @collection.internally_instrumented
307 def extend(self, items): ...
309 """
310 fn._sa_instrumented = True
311 return fn
313 @staticmethod
314 @util.deprecated(
315 "1.3",
316 "The :meth:`.collection.converter` handler is deprecated and will "
317 "be removed in a future release. Please refer to the "
318 ":class:`.AttributeEvents.bulk_replace` listener interface in "
319 "conjunction with the :func:`.event.listen` function.",
320 )
321 def converter(fn):
322 """Tag the method as the collection converter.
324 This optional method will be called when a collection is being
325 replaced entirely, as in::
327 myobj.acollection = [newvalue1, newvalue2]
329 The converter method will receive the object being assigned and should
330 return an iterable of values suitable for use by the ``appender``
331 method. A converter must not assign values or mutate the collection,
332 its sole job is to adapt the value the user provides into an iterable
333 of values for the ORM's use.
335 The default converter implementation will use duck-typing to do the
336 conversion. A dict-like collection will be convert into an iterable
337 of dictionary values, and other types will simply be iterated::
339 @collection.converter
340 def convert(self, other): ...
342 If the duck-typing of the object does not match the type of this
343 collection, a TypeError is raised.
345 Supply an implementation of this method if you want to expand the
346 range of possible types that can be assigned in bulk or perform
347 validation on the values about to be assigned.
349 """
350 fn._sa_instrument_role = "converter"
351 return fn
353 @staticmethod
354 def adds(arg):
355 """Mark the method as adding an entity to the collection.
357 Adds "add to collection" handling to the method. The decorator
358 argument indicates which method argument holds the SQLAlchemy-relevant
359 value. Arguments can be specified positionally (i.e. integer) or by
360 name::
362 @collection.adds(1)
363 def push(self, item): ...
365 @collection.adds('entity')
366 def do_stuff(self, thing, entity=None): ...
368 """
370 def decorator(fn):
371 fn._sa_instrument_before = ("fire_append_event", arg)
372 return fn
374 return decorator
376 @staticmethod
377 def replaces(arg):
378 """Mark the method as replacing an entity in the collection.
380 Adds "add to collection" and "remove from collection" handling to
381 the method. The decorator argument indicates which method argument
382 holds the SQLAlchemy-relevant value to be added, and return value, if
383 any will be considered the value to remove.
385 Arguments can be specified positionally (i.e. integer) or by name::
387 @collection.replaces(2)
388 def __setitem__(self, index, item): ...
390 """
392 def decorator(fn):
393 fn._sa_instrument_before = ("fire_append_event", arg)
394 fn._sa_instrument_after = "fire_remove_event"
395 return fn
397 return decorator
399 @staticmethod
400 def removes(arg):
401 """Mark the method as removing an entity in the collection.
403 Adds "remove from collection" handling to the method. The decorator
404 argument indicates which method argument holds the SQLAlchemy-relevant
405 value to be removed. Arguments can be specified positionally (i.e.
406 integer) or by name::
408 @collection.removes(1)
409 def zap(self, item): ...
411 For methods where the value to remove is not known at call-time, use
412 collection.removes_return.
414 """
416 def decorator(fn):
417 fn._sa_instrument_before = ("fire_remove_event", arg)
418 return fn
420 return decorator
422 @staticmethod
423 def removes_return():
424 """Mark the method as removing an entity in the collection.
426 Adds "remove from collection" handling to the method. The return
427 value of the method, if any, is considered the value to remove. The
428 method arguments are not inspected::
430 @collection.removes_return()
431 def pop(self): ...
433 For methods where the value to remove is known at call-time, use
434 collection.remove.
436 """
438 def decorator(fn):
439 fn._sa_instrument_after = "fire_remove_event"
440 return fn
442 return decorator
445if TYPE_CHECKING:
447 def collection_adapter(collection: Collection[Any]) -> CollectionAdapter:
448 """Fetch the :class:`.CollectionAdapter` for a collection."""
450else:
451 collection_adapter = operator.attrgetter("_sa_adapter")
454class CollectionAdapter:
455 """Bridges between the ORM and arbitrary Python collections.
457 Proxies base-level collection operations (append, remove, iterate)
458 to the underlying Python collection, and emits add/remove events for
459 entities entering or leaving the collection.
461 The ORM uses :class:`.CollectionAdapter` exclusively for interaction with
462 entity collections.
465 """
467 __slots__ = (
468 "attr",
469 "_key",
470 "_data",
471 "owner_state",
472 "_converter",
473 "invalidated",
474 "empty",
475 )
477 attr: CollectionAttributeImpl
478 _key: str
480 # this is actually a weakref; see note in constructor
481 _data: Callable[..., _AdaptedCollectionProtocol]
483 owner_state: InstanceState[Any]
484 _converter: _CollectionConverterProtocol
485 invalidated: bool
486 empty: bool
488 def __init__(
489 self,
490 attr: CollectionAttributeImpl,
491 owner_state: InstanceState[Any],
492 data: _AdaptedCollectionProtocol,
493 ):
494 self.attr = attr
495 self._key = attr.key
497 # this weakref stays referenced throughout the lifespan of
498 # CollectionAdapter. so while the weakref can return None, this
499 # is realistically only during garbage collection of this object, so
500 # we type this as a callable that returns _AdaptedCollectionProtocol
501 # in all cases.
502 self._data = weakref.ref(data) # type: ignore
504 self.owner_state = owner_state
505 data._sa_adapter = self
506 self._converter = data._sa_converter
507 self.invalidated = False
508 self.empty = False
510 def _warn_invalidated(self) -> None:
511 util.warn("This collection has been invalidated.")
513 @property
514 def data(self) -> _AdaptedCollectionProtocol:
515 "The entity collection being adapted."
516 return self._data()
518 @property
519 def _referenced_by_owner(self) -> bool:
520 """return True if the owner state still refers to this collection.
522 This will return False within a bulk replace operation,
523 where this collection is the one being replaced.
525 """
526 return self.owner_state.dict[self._key] is self._data()
528 def bulk_appender(self):
529 return self._data()._sa_appender
531 def append_with_event(
532 self, item: Any, initiator: Optional[AttributeEventToken] = None
533 ) -> None:
534 """Add an entity to the collection, firing mutation events."""
536 self._data()._sa_appender(item, _sa_initiator=initiator)
538 def _set_empty(self, user_data):
539 assert (
540 not self.empty
541 ), "This collection adapter is already in the 'empty' state"
542 self.empty = True
543 self.owner_state._empty_collections[self._key] = user_data
545 def _reset_empty(self) -> None:
546 assert (
547 self.empty
548 ), "This collection adapter is not in the 'empty' state"
549 self.empty = False
550 self.owner_state.dict[self._key] = (
551 self.owner_state._empty_collections.pop(self._key)
552 )
554 def _refuse_empty(self) -> NoReturn:
555 raise sa_exc.InvalidRequestError(
556 "This is a special 'empty' collection which cannot accommodate "
557 "internal mutation operations"
558 )
560 def append_without_event(self, item: Any) -> None:
561 """Add or restore an entity to the collection, firing no events."""
563 if self.empty:
564 self._refuse_empty()
565 self._data()._sa_appender(item, _sa_initiator=False)
567 def append_multiple_without_event(self, items: Iterable[Any]) -> None:
568 """Add or restore an entity to the collection, firing no events."""
569 if self.empty:
570 self._refuse_empty()
571 appender = self._data()._sa_appender
572 for item in items:
573 appender(item, _sa_initiator=False)
575 def bulk_remover(self):
576 return self._data()._sa_remover
578 def remove_with_event(
579 self, item: Any, initiator: Optional[AttributeEventToken] = None
580 ) -> None:
581 """Remove an entity from the collection, firing mutation events."""
582 self._data()._sa_remover(item, _sa_initiator=initiator)
584 def remove_without_event(self, item: Any) -> None:
585 """Remove an entity from the collection, firing no events."""
586 if self.empty:
587 self._refuse_empty()
588 self._data()._sa_remover(item, _sa_initiator=False)
590 def clear_with_event(
591 self, initiator: Optional[AttributeEventToken] = None
592 ) -> None:
593 """Empty the collection, firing a mutation event for each entity."""
595 if self.empty:
596 self._refuse_empty()
597 remover = self._data()._sa_remover
598 for item in list(self):
599 remover(item, _sa_initiator=initiator)
601 def clear_without_event(self) -> None:
602 """Empty the collection, firing no events."""
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=False)
610 def __iter__(self):
611 """Iterate over entities in the collection."""
613 return iter(self._data()._sa_iterator())
615 def __len__(self):
616 """Count entities in the collection."""
617 return len(list(self._data()._sa_iterator()))
619 def __bool__(self):
620 return True
622 def _fire_append_wo_mutation_event_bulk(
623 self, items, initiator=None, key=NO_KEY
624 ):
625 if not items:
626 return
628 if initiator is not False:
629 if self.invalidated:
630 self._warn_invalidated()
632 if self.empty:
633 self._reset_empty()
635 for item in items:
636 self.attr.fire_append_wo_mutation_event(
637 self.owner_state,
638 self.owner_state.dict,
639 item,
640 initiator,
641 key,
642 )
644 def fire_append_wo_mutation_event(self, item, initiator=None, key=NO_KEY):
645 """Notify that a entity is entering the collection but is already
646 present.
649 Initiator is a token owned by the InstrumentedAttribute that
650 initiated the membership mutation, and should be left as None
651 unless you are passing along an initiator value from a chained
652 operation.
654 .. versionadded:: 1.4.15
656 """
657 if initiator is not False:
658 if self.invalidated:
659 self._warn_invalidated()
661 if self.empty:
662 self._reset_empty()
664 return self.attr.fire_append_wo_mutation_event(
665 self.owner_state, self.owner_state.dict, item, initiator, key
666 )
667 else:
668 return item
670 def fire_append_event(self, item, initiator=None, key=NO_KEY):
671 """Notify that a entity has entered the collection.
673 Initiator is a token owned by the InstrumentedAttribute that
674 initiated the membership mutation, and should be left as None
675 unless you are passing along an initiator value from a chained
676 operation.
678 """
679 if initiator is not False:
680 if self.invalidated:
681 self._warn_invalidated()
683 if self.empty:
684 self._reset_empty()
686 return self.attr.fire_append_event(
687 self.owner_state, self.owner_state.dict, item, initiator, key
688 )
689 else:
690 return item
692 def _fire_remove_event_bulk(self, items, initiator=None, key=NO_KEY):
693 if not items:
694 return
696 if initiator is not False:
697 if self.invalidated:
698 self._warn_invalidated()
700 if self.empty:
701 self._reset_empty()
703 for item in items:
704 self.attr.fire_remove_event(
705 self.owner_state,
706 self.owner_state.dict,
707 item,
708 initiator,
709 key,
710 )
712 def fire_remove_event(self, item, initiator=None, key=NO_KEY):
713 """Notify that a entity has been removed from the collection.
715 Initiator is the InstrumentedAttribute that initiated the membership
716 mutation, and should be left as None unless you are passing along
717 an initiator value from a chained operation.
719 """
720 if initiator is not False:
721 if self.invalidated:
722 self._warn_invalidated()
724 if self.empty:
725 self._reset_empty()
727 self.attr.fire_remove_event(
728 self.owner_state, self.owner_state.dict, item, initiator, key
729 )
731 def fire_pre_remove_event(self, initiator=None, key=NO_KEY):
732 """Notify that an entity is about to be removed from the collection.
734 Only called if the entity cannot be removed after calling
735 fire_remove_event().
737 """
738 if self.invalidated:
739 self._warn_invalidated()
740 self.attr.fire_pre_remove_event(
741 self.owner_state,
742 self.owner_state.dict,
743 initiator=initiator,
744 key=key,
745 )
747 def __getstate__(self):
748 return {
749 "key": self._key,
750 "owner_state": self.owner_state,
751 "owner_cls": self.owner_state.class_,
752 "data": self.data,
753 "invalidated": self.invalidated,
754 "empty": self.empty,
755 }
757 def __setstate__(self, d):
758 self._key = d["key"]
759 self.owner_state = d["owner_state"]
761 # see note in constructor regarding this type: ignore
762 self._data = weakref.ref(d["data"]) # type: ignore
764 self._converter = d["data"]._sa_converter
765 d["data"]._sa_adapter = self
766 self.invalidated = d["invalidated"]
767 self.attr = getattr(d["owner_cls"], self._key).impl
768 self.empty = d.get("empty", False)
771def bulk_replace(values, existing_adapter, new_adapter, initiator=None):
772 """Load a new collection, firing events based on prior like membership.
774 Appends instances in ``values`` onto the ``new_adapter``. Events will be
775 fired for any instance not present in the ``existing_adapter``. Any
776 instances in ``existing_adapter`` not present in ``values`` will have
777 remove events fired upon them.
779 :param values: An iterable of collection member instances
781 :param existing_adapter: A :class:`.CollectionAdapter` of
782 instances to be replaced
784 :param new_adapter: An empty :class:`.CollectionAdapter`
785 to load with ``values``
788 """
790 assert isinstance(values, list)
792 idset = util.IdentitySet
793 existing_idset = idset(existing_adapter or ())
794 constants = existing_idset.intersection(values or ())
795 additions = idset(values or ()).difference(constants)
796 removals = existing_idset.difference(constants)
798 appender = new_adapter.bulk_appender()
800 for member in values or ():
801 if member in additions:
802 appender(member, _sa_initiator=initiator)
803 elif member in constants:
804 appender(member, _sa_initiator=False)
806 if existing_adapter:
807 existing_adapter._fire_append_wo_mutation_event_bulk(
808 constants, initiator=initiator
809 )
810 existing_adapter._fire_remove_event_bulk(removals, initiator=initiator)
813def prepare_instrumentation(
814 factory: Union[Type[Collection[Any]], _CollectionFactoryType],
815) -> _CollectionFactoryType:
816 """Prepare a callable for future use as a collection class factory.
818 Given a collection class factory (either a type or no-arg callable),
819 return another factory that will produce compatible instances when
820 called.
822 This function is responsible for converting collection_class=list
823 into the run-time behavior of collection_class=InstrumentedList.
825 """
827 impl_factory: _CollectionFactoryType
829 # Convert a builtin to 'Instrumented*'
830 if factory in __canned_instrumentation:
831 impl_factory = __canned_instrumentation[factory]
832 else:
833 impl_factory = cast(_CollectionFactoryType, factory)
835 cls: Union[_CollectionFactoryType, Type[Collection[Any]]]
837 # Create a specimen
838 cls = type(impl_factory())
840 # Did factory callable return a builtin?
841 if cls in __canned_instrumentation:
842 # if so, just convert.
843 # in previous major releases, this codepath wasn't working and was
844 # not covered by tests. prior to that it supplied a "wrapper"
845 # function that would return the class, though the rationale for this
846 # case is not known
847 impl_factory = __canned_instrumentation[cls]
848 cls = type(impl_factory())
850 # Instrument the class if needed.
851 if __instrumentation_mutex.acquire():
852 try:
853 if getattr(cls, "_sa_instrumented", None) != id(cls):
854 _instrument_class(cls)
855 finally:
856 __instrumentation_mutex.release()
858 return impl_factory
861def _instrument_class(cls):
862 """Modify methods in a class and install instrumentation."""
864 # In the normal call flow, a request for any of the 3 basic collection
865 # types is transformed into one of our trivial subclasses
866 # (e.g. InstrumentedList). Catch anything else that sneaks in here...
867 if cls.__module__ == "__builtin__":
868 raise sa_exc.ArgumentError(
869 "Can not instrument a built-in type. Use a "
870 "subclass, even a trivial one."
871 )
873 roles, methods = _locate_roles_and_methods(cls)
875 _setup_canned_roles(cls, roles, methods)
877 _assert_required_roles(cls, roles, methods)
879 _set_collection_attributes(cls, roles, methods)
882def _locate_roles_and_methods(cls):
883 """search for _sa_instrument_role-decorated methods in
884 method resolution order, assign to roles.
886 """
888 roles: Dict[str, str] = {}
889 methods: Dict[str, Tuple[Optional[str], Optional[int], Optional[str]]] = {}
891 for supercls in cls.__mro__:
892 for name, method in vars(supercls).items():
893 if not callable(method):
894 continue
896 # note role declarations
897 if hasattr(method, "_sa_instrument_role"):
898 role = method._sa_instrument_role
899 assert role in (
900 "appender",
901 "remover",
902 "iterator",
903 "converter",
904 )
905 roles.setdefault(role, name)
907 # transfer instrumentation requests from decorated function
908 # to the combined queue
909 before: Optional[Tuple[str, int]] = None
910 after: Optional[str] = None
912 if hasattr(method, "_sa_instrument_before"):
913 op, argument = method._sa_instrument_before
914 assert op in ("fire_append_event", "fire_remove_event")
915 before = op, argument
916 if hasattr(method, "_sa_instrument_after"):
917 op = method._sa_instrument_after
918 assert op in ("fire_append_event", "fire_remove_event")
919 after = op
920 if before:
921 methods[name] = before + (after,)
922 elif after:
923 methods[name] = None, None, after
924 return roles, methods
927def _setup_canned_roles(cls, roles, methods):
928 """see if this class has "canned" roles based on a known
929 collection type (dict, set, list). Apply those roles
930 as needed to the "roles" dictionary, and also
931 prepare "decorator" methods
933 """
934 collection_type = util.duck_type_collection(cls)
935 if collection_type in __interfaces:
936 assert collection_type is not None
937 canned_roles, decorators = __interfaces[collection_type]
938 for role, name in canned_roles.items():
939 roles.setdefault(role, name)
941 # apply ABC auto-decoration to methods that need it
942 for method, decorator in decorators.items():
943 fn = getattr(cls, method, None)
944 if (
945 fn
946 and method not in methods
947 and not hasattr(fn, "_sa_instrumented")
948 ):
949 setattr(cls, method, decorator(fn))
952def _assert_required_roles(cls, roles, methods):
953 """ensure all roles are present, and apply implicit instrumentation if
954 needed
956 """
957 if "appender" not in roles or not hasattr(cls, roles["appender"]):
958 raise sa_exc.ArgumentError(
959 "Type %s must elect an appender method to be "
960 "a collection class" % cls.__name__
961 )
962 elif roles["appender"] not in methods and not hasattr(
963 getattr(cls, roles["appender"]), "_sa_instrumented"
964 ):
965 methods[roles["appender"]] = ("fire_append_event", 1, None)
967 if "remover" not in roles or not hasattr(cls, roles["remover"]):
968 raise sa_exc.ArgumentError(
969 "Type %s must elect a remover method to be "
970 "a collection class" % cls.__name__
971 )
972 elif roles["remover"] not in methods and not hasattr(
973 getattr(cls, roles["remover"]), "_sa_instrumented"
974 ):
975 methods[roles["remover"]] = ("fire_remove_event", 1, None)
977 if "iterator" not in roles or not hasattr(cls, roles["iterator"]):
978 raise sa_exc.ArgumentError(
979 "Type %s must elect an iterator method to be "
980 "a collection class" % cls.__name__
981 )
984def _set_collection_attributes(cls, roles, methods):
985 """apply ad-hoc instrumentation from decorators, class-level defaults
986 and implicit role declarations
988 """
989 for method_name, (before, argument, after) in methods.items():
990 setattr(
991 cls,
992 method_name,
993 _instrument_membership_mutator(
994 getattr(cls, method_name), before, argument, after
995 ),
996 )
997 # intern the role map
998 for role, method_name in roles.items():
999 setattr(cls, "_sa_%s" % role, getattr(cls, method_name))
1001 cls._sa_adapter = None
1003 if not hasattr(cls, "_sa_converter"):
1004 cls._sa_converter = None
1005 cls._sa_instrumented = id(cls)
1008def _instrument_membership_mutator(method, before, argument, after):
1009 """Route method args and/or return value through the collection
1010 adapter."""
1011 # This isn't smart enough to handle @adds(1) for 'def fn(self, (a, b))'
1012 if before:
1013 fn_args = list(
1014 util.flatten_iterator(inspect_getfullargspec(method)[0])
1015 )
1016 if isinstance(argument, int):
1017 pos_arg = argument
1018 named_arg = len(fn_args) > argument and fn_args[argument] or None
1019 else:
1020 if argument in fn_args:
1021 pos_arg = fn_args.index(argument)
1022 else:
1023 pos_arg = None
1024 named_arg = argument
1025 del fn_args
1027 def wrapper(*args, **kw):
1028 if before:
1029 if pos_arg is None:
1030 if named_arg not in kw:
1031 raise sa_exc.ArgumentError(
1032 "Missing argument %s" % argument
1033 )
1034 value = kw[named_arg]
1035 else:
1036 if len(args) > pos_arg:
1037 value = args[pos_arg]
1038 elif named_arg in kw:
1039 value = kw[named_arg]
1040 else:
1041 raise sa_exc.ArgumentError(
1042 "Missing argument %s" % argument
1043 )
1045 initiator = kw.pop("_sa_initiator", None)
1046 if initiator is False:
1047 executor = None
1048 else:
1049 executor = args[0]._sa_adapter
1051 if before and executor:
1052 getattr(executor, before)(value, initiator)
1054 if not after or not executor:
1055 return method(*args, **kw)
1056 else:
1057 res = method(*args, **kw)
1058 if res is not None:
1059 getattr(executor, after)(res, initiator)
1060 return res
1062 wrapper._sa_instrumented = True # type: ignore[attr-defined]
1063 if hasattr(method, "_sa_instrument_role"):
1064 wrapper._sa_instrument_role = method._sa_instrument_role # type: ignore[attr-defined] # noqa: E501
1065 wrapper.__name__ = method.__name__
1066 wrapper.__doc__ = method.__doc__
1067 return wrapper
1070def __set_wo_mutation(collection, item, _sa_initiator=None):
1071 """Run set wo mutation events.
1073 The collection is not mutated.
1075 """
1076 if _sa_initiator is not False:
1077 executor = collection._sa_adapter
1078 if executor:
1079 executor.fire_append_wo_mutation_event(
1080 item, _sa_initiator, key=None
1081 )
1084def __set(collection, item, _sa_initiator, key):
1085 """Run set events.
1087 This event always occurs before the collection is actually mutated.
1089 """
1091 if _sa_initiator is not False:
1092 executor = collection._sa_adapter
1093 if executor:
1094 item = executor.fire_append_event(item, _sa_initiator, key=key)
1095 return item
1098def __del(collection, item, _sa_initiator, key):
1099 """Run del events.
1101 This event occurs before the collection is actually mutated, *except*
1102 in the case of a pop operation, in which case it occurs afterwards.
1103 For pop operations, the __before_pop hook is called before the
1104 operation occurs.
1106 """
1107 if _sa_initiator is not False:
1108 executor = collection._sa_adapter
1109 if executor:
1110 executor.fire_remove_event(item, _sa_initiator, key=key)
1113def __before_pop(collection, _sa_initiator=None):
1114 """An event which occurs on a before a pop() operation occurs."""
1115 executor = collection._sa_adapter
1116 if executor:
1117 executor.fire_pre_remove_event(_sa_initiator)
1120def _list_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1121 """Tailored instrumentation wrappers for any list-like class."""
1123 def _tidy(fn):
1124 fn._sa_instrumented = True
1125 fn.__doc__ = getattr(list, fn.__name__).__doc__
1127 def append(fn):
1128 def append(self, item, _sa_initiator=None):
1129 item = __set(self, item, _sa_initiator, NO_KEY)
1130 fn(self, item)
1132 _tidy(append)
1133 return append
1135 def remove(fn):
1136 def remove(self, value, _sa_initiator=None):
1137 __del(self, value, _sa_initiator, NO_KEY)
1138 # testlib.pragma exempt:__eq__
1139 fn(self, value)
1141 _tidy(remove)
1142 return remove
1144 def insert(fn):
1145 def insert(self, index, value):
1146 value = __set(self, value, None, index)
1147 fn(self, index, value)
1149 _tidy(insert)
1150 return insert
1152 def __setitem__(fn):
1153 def __setitem__(self, index, value):
1154 if not isinstance(index, slice):
1155 existing = self[index]
1156 if existing is not None:
1157 __del(self, existing, None, index)
1158 value = __set(self, value, None, index)
1159 fn(self, index, value)
1160 else:
1161 # slice assignment requires __delitem__, insert, __len__
1162 step = index.step or 1
1163 start = index.start or 0
1164 if start < 0:
1165 start += len(self)
1166 if index.stop is not None:
1167 stop = index.stop
1168 else:
1169 stop = len(self)
1170 if stop < 0:
1171 stop += len(self)
1173 if step == 1:
1174 if value is self:
1175 return
1176 for i in range(start, stop, step):
1177 if len(self) > start:
1178 del self[start]
1180 for i, item in enumerate(value):
1181 self.insert(i + start, item)
1182 else:
1183 rng = list(range(start, stop, step))
1184 if len(value) != len(rng):
1185 raise ValueError(
1186 "attempt to assign sequence of size %s to "
1187 "extended slice of size %s"
1188 % (len(value), len(rng))
1189 )
1190 for i, item in zip(rng, value):
1191 self.__setitem__(i, item)
1193 _tidy(__setitem__)
1194 return __setitem__
1196 def __delitem__(fn):
1197 def __delitem__(self, index):
1198 if not isinstance(index, slice):
1199 item = self[index]
1200 __del(self, item, None, index)
1201 fn(self, index)
1202 else:
1203 # slice deletion requires __getslice__ and a slice-groking
1204 # __getitem__ for stepped deletion
1205 # note: not breaking this into atomic dels
1206 for item in self[index]:
1207 __del(self, item, None, index)
1208 fn(self, index)
1210 _tidy(__delitem__)
1211 return __delitem__
1213 def extend(fn):
1214 def extend(self, iterable):
1215 for value in list(iterable):
1216 self.append(value)
1218 _tidy(extend)
1219 return extend
1221 def __iadd__(fn):
1222 def __iadd__(self, iterable):
1223 # list.__iadd__ takes any iterable and seems to let TypeError
1224 # raise as-is instead of returning NotImplemented
1225 for value in list(iterable):
1226 self.append(value)
1227 return self
1229 _tidy(__iadd__)
1230 return __iadd__
1232 def pop(fn):
1233 def pop(self, index=-1):
1234 __before_pop(self)
1235 item = fn(self, index)
1236 __del(self, item, None, index)
1237 return item
1239 _tidy(pop)
1240 return pop
1242 def clear(fn):
1243 def clear(self, index=-1):
1244 for item in self:
1245 __del(self, item, None, index)
1246 fn(self)
1248 _tidy(clear)
1249 return clear
1251 # __imul__ : not wrapping this. all members of the collection are already
1252 # present, so no need to fire appends... wrapping it with an explicit
1253 # decorator is still possible, so events on *= can be had if they're
1254 # desired. hard to imagine a use case for __imul__, though.
1256 l = locals().copy()
1257 l.pop("_tidy")
1258 return l
1261def _dict_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1262 """Tailored instrumentation wrappers for any dict-like mapping class."""
1264 def _tidy(fn):
1265 fn._sa_instrumented = True
1266 fn.__doc__ = getattr(dict, fn.__name__).__doc__
1268 def __setitem__(fn):
1269 def __setitem__(self, key, value, _sa_initiator=None):
1270 if key in self:
1271 __del(self, self[key], _sa_initiator, key)
1272 value = __set(self, value, _sa_initiator, key)
1273 fn(self, key, value)
1275 _tidy(__setitem__)
1276 return __setitem__
1278 def __delitem__(fn):
1279 def __delitem__(self, key, _sa_initiator=None):
1280 if key in self:
1281 __del(self, self[key], _sa_initiator, key)
1282 fn(self, key)
1284 _tidy(__delitem__)
1285 return __delitem__
1287 def clear(fn):
1288 def clear(self):
1289 for key in self:
1290 __del(self, self[key], None, key)
1291 fn(self)
1293 _tidy(clear)
1294 return clear
1296 def pop(fn):
1297 def pop(self, key, default=NO_ARG):
1298 __before_pop(self)
1299 _to_del = key in self
1300 if default is NO_ARG:
1301 item = fn(self, key)
1302 else:
1303 item = fn(self, key, default)
1304 if _to_del:
1305 __del(self, item, None, key)
1306 return item
1308 _tidy(pop)
1309 return pop
1311 def popitem(fn):
1312 def popitem(self):
1313 __before_pop(self)
1314 item = fn(self)
1315 __del(self, item[1], None, 1)
1316 return item
1318 _tidy(popitem)
1319 return popitem
1321 def setdefault(fn):
1322 def setdefault(self, key, default=None):
1323 if key not in self:
1324 self.__setitem__(key, default)
1325 return default
1326 else:
1327 value = self.__getitem__(key)
1328 if value is default:
1329 __set_wo_mutation(self, value, None)
1331 return value
1333 _tidy(setdefault)
1334 return setdefault
1336 def update(fn):
1337 def update(self, __other=NO_ARG, **kw):
1338 if __other is not NO_ARG:
1339 if hasattr(__other, "keys"):
1340 for key in list(__other):
1341 if key not in self or self[key] is not __other[key]:
1342 self[key] = __other[key]
1343 else:
1344 __set_wo_mutation(self, __other[key], None)
1345 else:
1346 for key, value in __other:
1347 if key not in self or self[key] is not value:
1348 self[key] = value
1349 else:
1350 __set_wo_mutation(self, value, None)
1351 for key in kw:
1352 if key not in self or self[key] is not kw[key]:
1353 self[key] = kw[key]
1354 else:
1355 __set_wo_mutation(self, kw[key], None)
1357 _tidy(update)
1358 return update
1360 l = locals().copy()
1361 l.pop("_tidy")
1362 return l
1365_set_binop_bases = (set, frozenset)
1368def _set_binops_check_strict(self: Any, obj: Any) -> bool:
1369 """Allow only set, frozenset and self.__class__-derived
1370 objects in binops."""
1371 return isinstance(obj, _set_binop_bases + (self.__class__,))
1374def _set_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1375 """Tailored instrumentation wrappers for any set-like class."""
1377 def _tidy(fn):
1378 fn._sa_instrumented = True
1379 fn.__doc__ = getattr(set, fn.__name__).__doc__
1381 def add(fn):
1382 def add(self, value, _sa_initiator=None):
1383 if value not in self:
1384 value = __set(self, value, _sa_initiator, NO_KEY)
1385 else:
1386 __set_wo_mutation(self, value, _sa_initiator)
1387 # testlib.pragma exempt:__hash__
1388 fn(self, value)
1390 _tidy(add)
1391 return add
1393 def discard(fn):
1394 def discard(self, value, _sa_initiator=None):
1395 # testlib.pragma exempt:__hash__
1396 if value in self:
1397 __del(self, value, _sa_initiator, NO_KEY)
1398 # testlib.pragma exempt:__hash__
1399 fn(self, value)
1401 _tidy(discard)
1402 return discard
1404 def remove(fn):
1405 def remove(self, value, _sa_initiator=None):
1406 # testlib.pragma exempt:__hash__
1407 if value in self:
1408 __del(self, value, _sa_initiator, NO_KEY)
1409 # testlib.pragma exempt:__hash__
1410 fn(self, value)
1412 _tidy(remove)
1413 return remove
1415 def pop(fn):
1416 def pop(self):
1417 __before_pop(self)
1418 item = fn(self)
1419 # for set in particular, we have no way to access the item
1420 # that will be popped before pop is called.
1421 __del(self, item, None, NO_KEY)
1422 return item
1424 _tidy(pop)
1425 return pop
1427 def clear(fn):
1428 def clear(self):
1429 for item in list(self):
1430 self.remove(item)
1432 _tidy(clear)
1433 return clear
1435 def update(fn):
1436 def update(self, value):
1437 for item in value:
1438 self.add(item)
1440 _tidy(update)
1441 return update
1443 def __ior__(fn):
1444 def __ior__(self, value):
1445 if not _set_binops_check_strict(self, value):
1446 return NotImplemented
1447 for item in value:
1448 self.add(item)
1449 return self
1451 _tidy(__ior__)
1452 return __ior__
1454 def difference_update(fn):
1455 def difference_update(self, value):
1456 for item in value:
1457 self.discard(item)
1459 _tidy(difference_update)
1460 return difference_update
1462 def __isub__(fn):
1463 def __isub__(self, value):
1464 if not _set_binops_check_strict(self, value):
1465 return NotImplemented
1466 for item in value:
1467 self.discard(item)
1468 return self
1470 _tidy(__isub__)
1471 return __isub__
1473 def intersection_update(fn):
1474 def intersection_update(self, other):
1475 want, have = self.intersection(other), set(self)
1476 remove, add = have - want, want - have
1478 for item in remove:
1479 self.remove(item)
1480 for item in add:
1481 self.add(item)
1483 _tidy(intersection_update)
1484 return intersection_update
1486 def __iand__(fn):
1487 def __iand__(self, other):
1488 if not _set_binops_check_strict(self, other):
1489 return NotImplemented
1490 want, have = self.intersection(other), set(self)
1491 remove, add = have - want, want - have
1493 for item in remove:
1494 self.remove(item)
1495 for item in add:
1496 self.add(item)
1497 return self
1499 _tidy(__iand__)
1500 return __iand__
1502 def symmetric_difference_update(fn):
1503 def symmetric_difference_update(self, other):
1504 want, have = self.symmetric_difference(other), set(self)
1505 remove, add = have - want, want - have
1507 for item in remove:
1508 self.remove(item)
1509 for item in add:
1510 self.add(item)
1512 _tidy(symmetric_difference_update)
1513 return symmetric_difference_update
1515 def __ixor__(fn):
1516 def __ixor__(self, other):
1517 if not _set_binops_check_strict(self, other):
1518 return NotImplemented
1519 want, have = self.symmetric_difference(other), set(self)
1520 remove, add = have - want, want - have
1522 for item in remove:
1523 self.remove(item)
1524 for item in add:
1525 self.add(item)
1526 return self
1528 _tidy(__ixor__)
1529 return __ixor__
1531 l = locals().copy()
1532 l.pop("_tidy")
1533 return l
1536class InstrumentedList(List[_T]):
1537 """An instrumented version of the built-in list."""
1540class InstrumentedSet(Set[_T]):
1541 """An instrumented version of the built-in set."""
1544class InstrumentedDict(Dict[_KT, _VT]):
1545 """An instrumented version of the built-in dict."""
1548__canned_instrumentation = cast(
1549 util.immutabledict[Any, _CollectionFactoryType],
1550 util.immutabledict(
1551 {
1552 list: InstrumentedList,
1553 set: InstrumentedSet,
1554 dict: InstrumentedDict,
1555 }
1556 ),
1557)
1559__interfaces: util.immutabledict[
1560 Any,
1561 Tuple[
1562 Dict[str, str],
1563 Dict[str, Callable[..., Any]],
1564 ],
1565] = util.immutabledict(
1566 {
1567 list: (
1568 {
1569 "appender": "append",
1570 "remover": "remove",
1571 "iterator": "__iter__",
1572 },
1573 _list_decorators(),
1574 ),
1575 set: (
1576 {"appender": "add", "remover": "remove", "iterator": "__iter__"},
1577 _set_decorators(),
1578 ),
1579 # decorators are required for dicts and object collections.
1580 dict: ({"iterator": "values"}, _dict_decorators()),
1581 }
1582)
1585def __go(lcls):
1586 global keyfunc_mapping, mapped_collection
1587 global column_keyed_dict, column_mapped_collection
1588 global MappedCollection, KeyFuncDict
1589 global attribute_keyed_dict, attribute_mapped_collection
1591 from .mapped_collection import keyfunc_mapping
1592 from .mapped_collection import column_keyed_dict
1593 from .mapped_collection import attribute_keyed_dict
1594 from .mapped_collection import KeyFuncDict
1596 from .mapped_collection import mapped_collection
1597 from .mapped_collection import column_mapped_collection
1598 from .mapped_collection import attribute_mapped_collection
1599 from .mapped_collection import MappedCollection
1601 # ensure instrumentation is associated with
1602 # these built-in classes; if a user-defined class
1603 # subclasses these and uses @internally_instrumented,
1604 # the superclass is otherwise not instrumented.
1605 # see [ticket:2406].
1606 _instrument_class(InstrumentedList)
1607 _instrument_class(InstrumentedSet)
1608 _instrument_class(KeyFuncDict)
1611__go(locals())