codekingpro/portable-devtools
114k
1# orm/collections.py
2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7# mypy: allow-untyped-defs, allow-untyped-calls
8
9"""Support for collections of mapped entities.
10
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.
15
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::
22
23 from sqlalchemy.orm.collections import collection
24
25
26 class MyClass:
27 # ...
28
29 @collection.adds(1)
30 def store(self, item):
31 self.data.append(item)
32
33 @collection.removes_return()
34 def pop(self):
35 return self.data.pop()
36
37The second approach is a bundle of targeted decorators that wrap appropriate
38append and remove notifiers around the mutation methods present in the
39standard Python ``list``, ``set`` and ``dict`` interfaces. These could be
40specified in terms of generic decorator recipes, but are instead hand-tooled
41for increased efficiency. The targeted decorators occasionally implement
42adapter-like behavior, such as mapping bulk-set methods (``extend``,
43``update``, ``__setslice__``, etc.) into the series of atomic mutation events
44that the ORM requires.
45
46The targeted decorators are used internally for automatic instrumentation of
47entity collection classes. Every collection class goes through a
48transformation process roughly like so:
49
501. If the class is a built-in, substitute a trivial sub-class
512. Is this class already instrumented?
523. Add in generic decorators
534. Sniff out the collection interface through duck-typing
545. Add targeted decoration to any undecorated interface method
55
56This process modifies the class at runtime, decorating methods and adding some
57bookkeeping properties. This isn't possible (or desirable) for built-in
58classes like ``list``, so trivial sub-classes are substituted to hold
59decoration::
60
61 class InstrumentedList(list):
62 pass
63
64Collection classes can be specified in ``relationship(collection_class=)`` as
65types or a function that returns an instance. Collection classes are
66inspected and instrumented during the mapper compilation phase. The
67collection_class callable will be executed once to produce a specimen
68instance, and the type of that specimen will be instrumented. Functions that
69return built-in types like ``lists`` will be adapted to produce instrumented
70instances.
71
72When extending a known type like ``list``, additional decorations are not
73generally not needed. Odds are, the extension method will delegate to a
74method that's already instrumented. For example::
75
76 class QueueIsh(list):
77 def push(self, item):
78 self.append(item)
79
80 def shift(self):
81 return self.pop(0)
82
83There's no need to decorate these methods. ``append`` and ``pop`` are already
84instrumented as part of the ``list`` interface. Decorating them would fire
85duplicate events, which should be avoided.
86
87The targeted decoration tries not to rely on other methods in the underlying
88collection class, but some are unavoidable. Many depend on 'read' methods
89being present to properly instrument a 'write', for example, ``__setitem__``
90needs ``__getitem__``. "Bulk" methods like ``update`` and ``extend`` may also
91reimplemented in terms of atomic appends and removes, so the ``extend``
92decoration will actually perform many ``append`` operations and not call the
93underlying method at all.
94
95Tight control over bulk operation and the firing of events is also possible by
96implementing the instrumentation internally in your methods. The basic
97instrumentation package works under the general assumption that collection
98mutation will not raise unusual exceptions. If you want to closely
99orchestrate append and remove events with exception management, internal
100instrumentation may be the answer. Within your method,
101``collection_adapter(self)`` will retrieve an object that you can use for
102explicit control over triggering append and remove events.
103
104The owning object and :class:`.CollectionAttributeImpl` are also reachable
105through the adapter, allowing for some very sophisticated behavior.
106
107"""
108from __future__ import annotations
109
110import operator
111import threading
112import typing
113from typing import Any
114from typing import Callable
115from typing import cast
116from typing import Collection
117from typing import Dict
118from typing import Iterable
119from typing import List
120from typing import NoReturn
121from typing import Optional
122from typing import Set
123from typing import Tuple
124from typing import Type
125from typing import TYPE_CHECKING
126from typing import TypeVar
127from typing import Union
128import weakref
129
130from .base import NO_KEY
131from .. import exc as sa_exc
132from .. import util
133from ..sql.base import NO_ARG
134from ..util.compat import inspect_getfullargspec
135from ..util.typing import Protocol
136
137if typing.TYPE_CHECKING:
138 from .attributes import AttributeEventToken
139 from .attributes import CollectionAttributeImpl
140 from .mapped_collection import attribute_keyed_dict
141 from .mapped_collection import column_keyed_dict
142 from .mapped_collection import keyfunc_mapping
143 from .mapped_collection import KeyFuncDict # noqa: F401
144 from .state import InstanceState
145
146
147__all__ = [
148 "collection",
149 "collection_adapter",
150 "keyfunc_mapping",
151 "column_keyed_dict",
152 "attribute_keyed_dict",
153 "KeyFuncDict",
154 # old names in < 2.0
155 "mapped_collection",
156 "column_mapped_collection",
157 "attribute_mapped_collection",
158 "MappedCollection",
159]
160
161__instrumentation_mutex = threading.Lock()
162
163
164_CollectionFactoryType = Callable[[], "_AdaptedCollectionProtocol"]
165
166_T = TypeVar("_T", bound=Any)
167_KT = TypeVar("_KT", bound=Any)
168_VT = TypeVar("_VT", bound=Any)
169_COL = TypeVar("_COL", bound="Collection[Any]")
170_FN = TypeVar("_FN", bound="Callable[..., Any]")
171
172
173class _CollectionConverterProtocol(Protocol):
174 def __call__(self, collection: _COL) -> _COL: ...
175
176
177class _AdaptedCollectionProtocol(Protocol):
178 _sa_adapter: CollectionAdapter
179 _sa_appender: Callable[..., Any]
180 _sa_remover: Callable[..., Any]
181 _sa_iterator: Callable[..., Iterable[Any]]
182 _sa_converter: _CollectionConverterProtocol
183
184
185class collection:
186 """Decorators for entity collection classes.
187
188 The decorators fall into two groups: annotations and interception recipes.
189
190 The annotating decorators (appender, remover, iterator, converter,
191 internally_instrumented) indicate the method's purpose and take no
192 arguments. They are not written with parens::
193
194 @collection.appender
195 def append(self, append): ...
196
197 The recipe decorators all require parens, even those that take no
198 arguments::
199
200 @collection.adds("entity")
201 def insert(self, position, entity): ...
202
203
204 @collection.removes_return()
205 def popitem(self): ...
206
207 """
208
209 # Bundled as a class solely for ease of use: packaging, doc strings,
210 # importability.
211
212 @staticmethod
213 def appender(fn):
214 """Tag the method as the collection appender.
215
216 The appender method is called with one positional argument: the value
217 to append. The method will be automatically decorated with 'adds(1)'
218 if not already decorated::
219
220 @collection.appender
221 def add(self, append): ...
222
223
224 # or, equivalently
225 @collection.appender
226 @collection.adds(1)
227 def add(self, append): ...
228
229
230 # for mapping type, an 'append' may kick out a previous value
231 # that occupies that slot. consider d['a'] = 'foo'- any previous
232 # value in d['a'] is discarded.
233 @collection.appender
234 @collection.replaces(1)
235 def add(self, entity):
236 key = some_key_func(entity)
237 previous = None
238 if key in self:
239 previous = self[key]
240 self[key] = entity
241 return previous
242
243 If the value to append is not allowed in the collection, you may
244 raise an exception. Something to remember is that the appender
245 will be called for each object mapped by a database query. If the
246 database contains rows that violate your collection semantics, you
247 will need to get creative to fix the problem, as access via the
248 collection will not work.
249
250 If the appender method is internally instrumented, you must also
251 receive the keyword argument '_sa_initiator' and ensure its
252 promulgation to collection events.
253
254 """
255 fn._sa_instrument_role = "appender"
256 return fn
257
258 @staticmethod
259 def remover(fn):
260 """Tag the method as the collection remover.
261
262 The remover method is called with one positional argument: the value
263 to remove. The method will be automatically decorated with
264 :meth:`removes_return` if not already decorated::
265
266 @collection.remover
267 def zap(self, entity): ...
268
269
270 # or, equivalently
271 @collection.remover
272 @collection.removes_return()
273 def zap(self): ...
274
275 If the value to remove is not present in the collection, you may
276 raise an exception or return None to ignore the error.
277
278 If the remove method is internally instrumented, you must also
279 receive the keyword argument '_sa_initiator' and ensure its
280 promulgation to collection events.
281
282 """
283 fn._sa_instrument_role = "remover"
284 return fn
285
286 @staticmethod
287 def iterator(fn):
288 """Tag the method as the collection remover.
289
290 The iterator method is called with no arguments. It is expected to
291 return an iterator over all collection members::
292
293 @collection.iterator
294 def __iter__(self): ...
295
296 """
297 fn._sa_instrument_role = "iterator"
298 return fn
299
300 @staticmethod
301 def internally_instrumented(fn):
302 """Tag the method as instrumented.
303
304 This tag will prevent any decoration from being applied to the
305 method. Use this if you are orchestrating your own calls to
306 :func:`.collection_adapter` in one of the basic SQLAlchemy
307 interface methods, or to prevent an automatic ABC method
308 decoration from wrapping your implementation::
309
310 # normally an 'extend' method on a list-like class would be
311 # automatically intercepted and re-implemented in terms of
312 # SQLAlchemy events and append(). your implementation will
313 # never be called, unless:
314 @collection.internally_instrumented
315 def extend(self, items): ...
316
317 """
318 fn._sa_instrumented = True
319 return fn
320
321 @staticmethod
322 @util.deprecated(
323 "1.3",
324 "The :meth:`.collection.converter` handler is deprecated and will "
325 "be removed in a future release. Please refer to the "
326 ":class:`.AttributeEvents.bulk_replace` listener interface in "
327 "conjunction with the :func:`.event.listen` function.",
328 )
329 def converter(fn):
330 """Tag the method as the collection converter.
331
332 This optional method will be called when a collection is being
333 replaced entirely, as in::
334
335 myobj.acollection = [newvalue1, newvalue2]
336
337 The converter method will receive the object being assigned and should
338 return an iterable of values suitable for use by the ``appender``
339 method. A converter must not assign values or mutate the collection,
340 its sole job is to adapt the value the user provides into an iterable
341 of values for the ORM's use.
342
343 The default converter implementation will use duck-typing to do the
344 conversion. A dict-like collection will be convert into an iterable
345 of dictionary values, and other types will simply be iterated::
346
347 @collection.converter
348 def convert(self, other): ...
349
350 If the duck-typing of the object does not match the type of this
351 collection, a TypeError is raised.
352
353 Supply an implementation of this method if you want to expand the
354 range of possible types that can be assigned in bulk or perform
355 validation on the values about to be assigned.
356
357 """
358 fn._sa_instrument_role = "converter"
359 return fn
360
361 @staticmethod
362 def adds(arg: int) -> Callable[[_FN], _FN]:
363 """Mark the method as adding an entity to the collection.
364
365 Adds "add to collection" handling to the method. The decorator
366 argument indicates which method argument holds the SQLAlchemy-relevant
367 value. Arguments can be specified positionally (i.e. integer) or by
368 name::
369
370 @collection.adds(1)
371 def push(self, item): ...
372
373
374 @collection.adds("entity")
375 def do_stuff(self, thing, entity=None): ...
376
377 """
378
379 def decorator(fn):
380 fn._sa_instrument_before = ("fire_append_event", arg)
381 return fn
382
383 return decorator
384
385 @staticmethod
386 def replaces(arg):
387 """Mark the method as replacing an entity in the collection.
388
389 Adds "add to collection" and "remove from collection" handling to
390 the method. The decorator argument indicates which method argument
391 holds the SQLAlchemy-relevant value to be added, and return value, if
392 any will be considered the value to remove.
393
394 Arguments can be specified positionally (i.e. integer) or by name::
395
396 @collection.replaces(2)
397 def __setitem__(self, index, item): ...
398
399 """
400
401 def decorator(fn):
402 fn._sa_instrument_before = ("fire_append_event", arg)
403 fn._sa_instrument_after = "fire_remove_event"
404 return fn
405
406 return decorator
407
408 @staticmethod
409 def removes(arg):
410 """Mark the method as removing an entity in the collection.
411
412 Adds "remove from collection" handling to the method. The decorator
413 argument indicates which method argument holds the SQLAlchemy-relevant
414 value to be removed. Arguments can be specified positionally (i.e.
415 integer) or by name::
416
417 @collection.removes(1)
418 def zap(self, item): ...
419
420 For methods where the value to remove is not known at call-time, use
421 collection.removes_return.
422
423 """
424
425 def decorator(fn):
426 fn._sa_instrument_before = ("fire_remove_event", arg)
427 return fn
428
429 return decorator
430
431 @staticmethod
432 def removes_return():
433 """Mark the method as removing an entity in the collection.
434
435 Adds "remove from collection" handling to the method. The return
436 value of the method, if any, is considered the value to remove. The
437 method arguments are not inspected::
438
439 @collection.removes_return()
440 def pop(self): ...
441
442 For methods where the value to remove is known at call-time, use
443 collection.remove.
444
445 """
446
447 def decorator(fn):
448 fn._sa_instrument_after = "fire_remove_event"
449 return fn
450
451 return decorator
452
453
454if TYPE_CHECKING:
455
456 def collection_adapter(collection: Collection[Any]) -> CollectionAdapter:
457 """Fetch the :class:`.CollectionAdapter` for a collection."""
458
459else:
460 collection_adapter = operator.attrgetter("_sa_adapter")
461
462
463class CollectionAdapter:
464 """Bridges between the ORM and arbitrary Python collections.
465
466 Proxies base-level collection operations (append, remove, iterate)
467 to the underlying Python collection, and emits add/remove events for
468 entities entering or leaving the collection.
469
470 The ORM uses :class:`.CollectionAdapter` exclusively for interaction with
471 entity collections.
472
473
474 """
475
476 __slots__ = (
477 "attr",
478 "_key",
479 "_data",
480 "owner_state",
481 "_converter",
482 "invalidated",
483 "empty",
484 )
485
486 attr: CollectionAttributeImpl
487 _key: str
488
489 # this is actually a weakref; see note in constructor
490 _data: Callable[..., _AdaptedCollectionProtocol]
491
492 owner_state: InstanceState[Any]
493 _converter: _CollectionConverterProtocol
494 invalidated: bool
495 empty: bool
496
497 def __init__(
498 self,
499 attr: CollectionAttributeImpl,
500 owner_state: InstanceState[Any],
501 data: _AdaptedCollectionProtocol,
502 ):
503 self.attr = attr
504 self._key = attr.key
505
506 # this weakref stays referenced throughout the lifespan of
507 # CollectionAdapter. so while the weakref can return None, this
508 # is realistically only during garbage collection of this object, so
509 # we type this as a callable that returns _AdaptedCollectionProtocol
510 # in all cases.
511 self._data = weakref.ref(data) # type: ignore
512
513 self.owner_state = owner_state
514 data._sa_adapter = self
515 self._converter = data._sa_converter
516 self.invalidated = False
517 self.empty = False
518
519 def _warn_invalidated(self) -> None:
520 util.warn("This collection has been invalidated.")
521
522 @property
523 def data(self) -> _AdaptedCollectionProtocol:
524 "The entity collection being adapted."
525 return self._data()
526
527 @property
528 def _referenced_by_owner(self) -> bool:
529 """return True if the owner state still refers to this collection.
530
531 This will return False within a bulk replace operation,
532 where this collection is the one being replaced.
533
534 """
535 return self.owner_state.dict[self._key] is self._data()
536
537 def bulk_appender(self):
538 return self._data()._sa_appender
539
540 def append_with_event(
541 self, item: Any, initiator: Optional[AttributeEventToken] = None
542 ) -> None:
543 """Add an entity to the collection, firing mutation events."""
544
545 self._data()._sa_appender(item, _sa_initiator=initiator)
546
547 def _set_empty(self, user_data):
548 assert (
549 not self.empty
550 ), "This collection adapter is already in the 'empty' state"
551 self.empty = True
552 self.owner_state._empty_collections[self._key] = user_data
553
554 def _reset_empty(self) -> None:
555 assert (
556 self.empty
557 ), "This collection adapter is not in the 'empty' state"
558 self.empty = False
559 self.owner_state.dict[self._key] = (
560 self.owner_state._empty_collections.pop(self._key)
561 )
562
563 def _refuse_empty(self) -> NoReturn:
564 raise sa_exc.InvalidRequestError(
565 "This is a special 'empty' collection which cannot accommodate "
566 "internal mutation operations"
567 )
568
569 def append_without_event(self, item: Any) -> None:
570 """Add or restore an entity to the collection, firing no events."""
571
572 if self.empty:
573 self._refuse_empty()
574 self._data()._sa_appender(item, _sa_initiator=False)
575
576 def append_multiple_without_event(self, items: Iterable[Any]) -> None:
577 """Add or restore an entity to the collection, firing no events."""
578 if self.empty:
579 self._refuse_empty()
580 appender = self._data()._sa_appender
581 for item in items:
582 appender(item, _sa_initiator=False)
583
584 def bulk_remover(self):
585 return self._data()._sa_remover
586
587 def remove_with_event(
588 self, item: Any, initiator: Optional[AttributeEventToken] = None
589 ) -> None:
590 """Remove an entity from the collection, firing mutation events."""
591 self._data()._sa_remover(item, _sa_initiator=initiator)
592
593 def remove_without_event(self, item: Any) -> None:
594 """Remove an entity from the collection, firing no events."""
595 if self.empty:
596 self._refuse_empty()
597 self._data()._sa_remover(item, _sa_initiator=False)
598
599 def clear_with_event(
600 self, initiator: Optional[AttributeEventToken] = None
601 ) -> None:
602 """Empty the collection, firing a mutation event for each entity."""
603
604 if self.empty:
605 self._refuse_empty()
606 remover = self._data()._sa_remover
607 for item in list(self):
608 remover(item, _sa_initiator=initiator)
609
610 def clear_without_event(self) -> None:
611 """Empty the collection, firing no events."""
612
613 if self.empty:
614 self._refuse_empty()
615 remover = self._data()._sa_remover
616 for item in list(self):
617 remover(item, _sa_initiator=False)
618
619 def __iter__(self):
620 """Iterate over entities in the collection."""
621
622 return iter(self._data()._sa_iterator())
623
624 def __len__(self):
625 """Count entities in the collection."""
626 return len(list(self._data()._sa_iterator()))
627
628 def __bool__(self):
629 return True
630
631 def _fire_append_wo_mutation_event_bulk(
632 self, items, initiator=None, key=NO_KEY
633 ):
634 if not items:
635 return
636
637 if initiator is not False:
638 if self.invalidated:
639 self._warn_invalidated()
640
641 if self.empty:
642 self._reset_empty()
643
644 for item in items:
645 self.attr.fire_append_wo_mutation_event(
646 self.owner_state,
647 self.owner_state.dict,
648 item,
649 initiator,
650 key,
651 )
652
653 def fire_append_wo_mutation_event(self, item, initiator=None, key=NO_KEY):
654 """Notify that a entity is entering the collection but is already
655 present.
656
657
658 Initiator is a token owned by the InstrumentedAttribute that
659 initiated the membership mutation, and should be left as None
660 unless you are passing along an initiator value from a chained
661 operation.
662
663 .. versionadded:: 1.4.15
664
665 """
666 if initiator is not False:
667 if self.invalidated:
668 self._warn_invalidated()
669
670 if self.empty:
671 self._reset_empty()
672
673 return self.attr.fire_append_wo_mutation_event(
674 self.owner_state, self.owner_state.dict, item, initiator, key
675 )
676 else:
677 return item
678
679 def fire_append_event(self, item, initiator=None, key=NO_KEY):
680 """Notify that a entity has entered the collection.
681
682 Initiator is a token owned by the InstrumentedAttribute that
683 initiated the membership mutation, and should be left as None
684 unless you are passing along an initiator value from a chained
685 operation.
686
687 """
688 if initiator is not False:
689 if self.invalidated:
690 self._warn_invalidated()
691
692 if self.empty:
693 self._reset_empty()
694
695 return self.attr.fire_append_event(
696 self.owner_state, self.owner_state.dict, item, initiator, key
697 )
698 else:
699 return item
700
701 def _fire_remove_event_bulk(self, items, initiator=None, key=NO_KEY):
702 if not items:
703 return
704
705 if initiator is not False:
706 if self.invalidated:
707 self._warn_invalidated()
708
709 if self.empty:
710 self._reset_empty()
711
712 for item in items:
713 self.attr.fire_remove_event(
714 self.owner_state,
715 self.owner_state.dict,
716 item,
717 initiator,
718 key,
719 )
720
721 def fire_remove_event(self, item, initiator=None, key=NO_KEY):
722 """Notify that a entity has been removed from the collection.
723
724 Initiator is the InstrumentedAttribute that initiated the membership
725 mutation, and should be left as None unless you are passing along
726 an initiator value from a chained operation.
727
728 """
729 if initiator is not False:
730 if self.invalidated:
731 self._warn_invalidated()
732
733 if self.empty:
734 self._reset_empty()
735
736 self.attr.fire_remove_event(
737 self.owner_state, self.owner_state.dict, item, initiator, key
738 )
739
740 def fire_pre_remove_event(self, initiator=None, key=NO_KEY):
741 """Notify that an entity is about to be removed from the collection.
742
743 Only called if the entity cannot be removed after calling
744 fire_remove_event().
745
746 """
747 if self.invalidated:
748 self._warn_invalidated()
749 self.attr.fire_pre_remove_event(
750 self.owner_state,
751 self.owner_state.dict,
752 initiator=initiator,
753 key=key,
754 )
755
756 def __getstate__(self):
757 return {
758 "key": self._key,
759 "owner_state": self.owner_state,
760 "owner_cls": self.owner_state.class_,
761 "data": self.data,
762 "invalidated": self.invalidated,
763 "empty": self.empty,
764 }
765
766 def __setstate__(self, d):
767 self._key = d["key"]
768 self.owner_state = d["owner_state"]
769
770 # see note in constructor regarding this type: ignore
771 self._data = weakref.ref(d["data"]) # type: ignore
772
773 self._converter = d["data"]._sa_converter
774 d["data"]._sa_adapter = self
775 self.invalidated = d["invalidated"]
776 self.attr = getattr(d["owner_cls"], self._key).impl
777 self.empty = d.get("empty", False)
778
779
780def bulk_replace(values, existing_adapter, new_adapter, initiator=None):
781 """Load a new collection, firing events based on prior like membership.
782
783 Appends instances in ``values`` onto the ``new_adapter``. Events will be
784 fired for any instance not present in the ``existing_adapter``. Any
785 instances in ``existing_adapter`` not present in ``values`` will have
786 remove events fired upon them.
787
788 :param values: An iterable of collection member instances
789
790 :param existing_adapter: A :class:`.CollectionAdapter` of
791 instances to be replaced
792
793 :param new_adapter: An empty :class:`.CollectionAdapter`
794 to load with ``values``
795
796
797 """
798
799 assert isinstance(values, list)
800
801 idset = util.IdentitySet
802 existing_idset = idset(existing_adapter or ())
803 constants = existing_idset.intersection(values or ())
804 additions = idset(values or ()).difference(constants)
805 removals = existing_idset.difference(constants)
806
807 appender = new_adapter.bulk_appender()
808
809 for member in values or ():
810 if member in additions:
811 appender(member, _sa_initiator=initiator)
812 elif member in constants:
813 appender(member, _sa_initiator=False)
814
815 if existing_adapter:
816 existing_adapter._fire_append_wo_mutation_event_bulk(
817 constants, initiator=initiator
818 )
819 existing_adapter._fire_remove_event_bulk(removals, initiator=initiator)
820
821
822def prepare_instrumentation(
823 factory: Union[Type[Collection[Any]], _CollectionFactoryType],
824) -> _CollectionFactoryType:
825 """Prepare a callable for future use as a collection class factory.
826
827 Given a collection class factory (either a type or no-arg callable),
828 return another factory that will produce compatible instances when
829 called.
830
831 This function is responsible for converting collection_class=list
832 into the run-time behavior of collection_class=InstrumentedList.
833
834 """
835
836 impl_factory: _CollectionFactoryType
837
838 # Convert a builtin to 'Instrumented*'
839 if factory in __canned_instrumentation:
840 impl_factory = __canned_instrumentation[factory]
841 else:
842 impl_factory = cast(_CollectionFactoryType, factory)
843
844 cls: Union[_CollectionFactoryType, Type[Collection[Any]]]
845
846 # Create a specimen
847 cls = type(impl_factory())
848
849 # Did factory callable return a builtin?
850 if cls in __canned_instrumentation:
851 # if so, just convert.
852 # in previous major releases, this codepath wasn't working and was
853 # not covered by tests. prior to that it supplied a "wrapper"
854 # function that would return the class, though the rationale for this
855 # case is not known
856 impl_factory = __canned_instrumentation[cls]
857 cls = type(impl_factory())
858
859 # Instrument the class if needed.
860 if __instrumentation_mutex.acquire():
861 try:
862 if getattr(cls, "_sa_instrumented", None) != id(cls):
863 _instrument_class(cls)
864 finally:
865 __instrumentation_mutex.release()
866
867 return impl_factory
868
869
870def _instrument_class(cls):
871 """Modify methods in a class and install instrumentation."""
872
873 # In the normal call flow, a request for any of the 3 basic collection
874 # types is transformed into one of our trivial subclasses
875 # (e.g. InstrumentedList). Catch anything else that sneaks in here...
876 if cls.__module__ == "__builtin__":
877 raise sa_exc.ArgumentError(
878 "Can not instrument a built-in type. Use a "
879 "subclass, even a trivial one."
880 )
881
882 roles, methods = _locate_roles_and_methods(cls)
883
884 _setup_canned_roles(cls, roles, methods)
885
886 _assert_required_roles(cls, roles, methods)
887
888 _set_collection_attributes(cls, roles, methods)
889
890
891def _locate_roles_and_methods(cls):
892 """search for _sa_instrument_role-decorated methods in
893 method resolution order, assign to roles.
894
895 """
896
897 roles: Dict[str, str] = {}
898 methods: Dict[str, Tuple[Optional[str], Optional[int], Optional[str]]] = {}
899
900 for supercls in cls.__mro__:
901 for name, method in vars(supercls).items():
902 if not callable(method):
903 continue
904
905 # note role declarations
906 if hasattr(method, "_sa_instrument_role"):
907 role = method._sa_instrument_role
908 assert role in (
909 "appender",
910 "remover",
911 "iterator",
912 "converter",
913 )
914 roles.setdefault(role, name)
915
916 # transfer instrumentation requests from decorated function
917 # to the combined queue
918 before: Optional[Tuple[str, int]] = None
919 after: Optional[str] = None
920
921 if hasattr(method, "_sa_instrument_before"):
922 op, argument = method._sa_instrument_before
923 assert op in ("fire_append_event", "fire_remove_event")
924 before = op, argument
925 if hasattr(method, "_sa_instrument_after"):
926 op = method._sa_instrument_after
927 assert op in ("fire_append_event", "fire_remove_event")
928 after = op
929 if before:
930 methods[name] = before + (after,)
931 elif after:
932 methods[name] = None, None, after
933 return roles, methods
934
935
936def _setup_canned_roles(cls, roles, methods):
937 """see if this class has "canned" roles based on a known
938 collection type (dict, set, list). Apply those roles
939 as needed to the "roles" dictionary, and also
940 prepare "decorator" methods
941
942 """
943 collection_type = util.duck_type_collection(cls)
944 if collection_type in __interfaces:
945 assert collection_type is not None
946 canned_roles, decorators = __interfaces[collection_type]
947 for role, name in canned_roles.items():
948 roles.setdefault(role, name)
949
950 # apply ABC auto-decoration to methods that need it
951 for method, decorator in decorators.items():
952 fn = getattr(cls, method, None)
953 if (
954 fn
955 and method not in methods
956 and not hasattr(fn, "_sa_instrumented")
957 ):
958 setattr(cls, method, decorator(fn))
959
960
961def _assert_required_roles(cls, roles, methods):
962 """ensure all roles are present, and apply implicit instrumentation if
963 needed
964
965 """
966 if "appender" not in roles or not hasattr(cls, roles["appender"]):
967 raise sa_exc.ArgumentError(
968 "Type %s must elect an appender method to be "
969 "a collection class" % cls.__name__
970 )
971 elif roles["appender"] not in methods and not hasattr(
972 getattr(cls, roles["appender"]), "_sa_instrumented"
973 ):
974 methods[roles["appender"]] = ("fire_append_event", 1, None)
975
976 if "remover" not in roles or not hasattr(cls, roles["remover"]):
977 raise sa_exc.ArgumentError(
978 "Type %s must elect a remover method to be "
979 "a collection class" % cls.__name__
980 )
981 elif roles["remover"] not in methods and not hasattr(
982 getattr(cls, roles["remover"]), "_sa_instrumented"
983 ):
984 methods[roles["remover"]] = ("fire_remove_event", 1, None)
985
986 if "iterator" not in roles or not hasattr(cls, roles["iterator"]):
987 raise sa_exc.ArgumentError(
988 "Type %s must elect an iterator method to be "
989 "a collection class" % cls.__name__
990 )
991
992
993def _set_collection_attributes(cls, roles, methods):
994 """apply ad-hoc instrumentation from decorators, class-level defaults
995 and implicit role declarations
996
997 """
998 for method_name, (before, argument, after) in methods.items():
999 setattr(
1000 cls,
1001 method_name,
1002 _instrument_membership_mutator(
1003 getattr(cls, method_name), before, argument, after
1004 ),
1005 )
1006 # intern the role map
1007 for role, method_name in roles.items():
1008 setattr(cls, "_sa_%s" % role, getattr(cls, method_name))
1009
1010 cls._sa_adapter = None
1011
1012 if not hasattr(cls, "_sa_converter"):
1013 cls._sa_converter = None
1014 cls._sa_instrumented = id(cls)
1015
1016
1017def _instrument_membership_mutator(method, before, argument, after):
1018 """Route method args and/or return value through the collection
1019 adapter."""
1020 # This isn't smart enough to handle @adds(1) for 'def fn(self, (a, b))'
1021 if before:
1022 fn_args = list(
1023 util.flatten_iterator(inspect_getfullargspec(method)[0])
1024 )
1025 if isinstance(argument, int):
1026 pos_arg = argument
1027 named_arg = len(fn_args) > argument and fn_args[argument] or None
1028 else:
1029 if argument in fn_args:
1030 pos_arg = fn_args.index(argument)
1031 else:
1032 pos_arg = None
1033 named_arg = argument
1034 del fn_args
1035
1036 def wrapper(*args, **kw):
1037 if before:
1038 if pos_arg is None:
1039 if named_arg not in kw:
1040 raise sa_exc.ArgumentError(
1041 "Missing argument %s" % argument
1042 )
1043 value = kw[named_arg]
1044 else:
1045 if len(args) > pos_arg:
1046 value = args[pos_arg]
1047 elif named_arg in kw:
1048 value = kw[named_arg]
1049 else:
1050 raise sa_exc.ArgumentError(
1051 "Missing argument %s" % argument
1052 )
1053
1054 initiator = kw.pop("_sa_initiator", None)
1055 if initiator is False:
1056 executor = None
1057 else:
1058 executor = args[0]._sa_adapter
1059
1060 if before and executor:
1061 getattr(executor, before)(value, initiator)
1062
1063 if not after or not executor:
1064 return method(*args, **kw)
1065 else:
1066 res = method(*args, **kw)
1067 if res is not None:
1068 getattr(executor, after)(res, initiator)
1069 return res
1070
1071 wrapper._sa_instrumented = True # type: ignore[attr-defined]
1072 if hasattr(method, "_sa_instrument_role"):
1073 wrapper._sa_instrument_role = method._sa_instrument_role # type: ignore[attr-defined] # noqa: E501
1074 wrapper.__name__ = method.__name__
1075 wrapper.__doc__ = method.__doc__
1076 return wrapper
1077
1078
1079def __set_wo_mutation(collection, item, _sa_initiator=None):
1080 """Run set wo mutation events.
1081
1082 The collection is not mutated.
1083
1084 """
1085 if _sa_initiator is not False:
1086 executor = collection._sa_adapter
1087 if executor:
1088 executor.fire_append_wo_mutation_event(
1089 item, _sa_initiator, key=None
1090 )
1091
1092
1093def __set(collection, item, _sa_initiator, key):
1094 """Run set events.
1095
1096 This event always occurs before the collection is actually mutated.
1097
1098 """
1099
1100 if _sa_initiator is not False:
1101 executor = collection._sa_adapter
1102 if executor:
1103 item = executor.fire_append_event(item, _sa_initiator, key=key)
1104 return item
1105
1106
1107def __del(collection, item, _sa_initiator, key):
1108 """Run del events.
1109
1110 This event occurs before the collection is actually mutated, *except*
1111 in the case of a pop operation, in which case it occurs afterwards.
1112 For pop operations, the __before_pop hook is called before the
1113 operation occurs.
1114
1115 """
1116 if _sa_initiator is not False:
1117 executor = collection._sa_adapter
1118 if executor:
1119 executor.fire_remove_event(item, _sa_initiator, key=key)
1120
1121
1122def __before_pop(collection, _sa_initiator=None):
1123 """An event which occurs on a before a pop() operation occurs."""
1124 executor = collection._sa_adapter
1125 if executor:
1126 executor.fire_pre_remove_event(_sa_initiator)
1127
1128
1129def _list_decorators() -> Dict[str, Callable[[_FN], _FN]]:
1130 """Tailored instrumentation wrappers for any list-like class."""
1131
1132 def _tidy(fn):
1133 fn._sa_instrumented = True
1134 fn.__doc__ = getattr(list, fn.__name__).__doc__
1135
1136 def append(fn):
1137 def append(self, item, _sa_initiator=None):
1138 item = __set(self, item, _sa_initiator, NO_KEY)
1139 fn(self, item)
1140
1141 _tidy(append)
1142 return append
1143
1144 def remove(fn):
1145 def remove(self, value, _sa_initiator=None):
1146 __del(self, value, _sa_initiator, NO_KEY)
1147 # testlib.pragma exempt:__eq__
1148 fn(self, value)
1149
1150 _tidy(remove)
1151 return remove
1152
1153 def insert(fn):
1154 def insert(self, index, value):
1155 value = __set(self, value, None, index)
1156 fn(self, index, value)
1157
1158 _tidy(insert)
1159 return insert
1160
1161 def __setitem__(fn):
1162 def __setitem__(self, index, value):
1163 if not isinstance(index, slice):
1164 existing = self[index]
1165 if existing is not None:
1166 __del(self, existing, None, index)
1167 value = __set(self, value, None, index)
1168 fn(self, index, value)
1169 else:
1170 # slice assignment requires __delitem__, insert, __len__
1171 step = index.step or 1
1172 start = index.start or 0
1173 if start < 0:
1174 start += len(self)
1175 if index.stop is not None:
1176 stop = index.stop
1177 else:
1178 stop = len(self)
1179 if stop < 0:
1180 stop += len(self)
1181
1182 if step == 1:
1183 if value is self:
1184 return
1185 for i in range(start, stop, step):
1186 if len(self) > start:
1187 del self[start]
1188
1189 for i, item in enumerate(value):
1190 self.insert(i + start, item)
1191 else:
1192 rng = list(range(start, stop, step))
1193 if len(value) != len(rng):
1194 raise ValueError(
1195 "attempt to assign sequence of size %s to "
1196 "extended slice of size %s"
1197 % (len(value), len(rng))
1198 )
1199 for i, item in zip(rng, value):
1200 self.__setitem__(i, item)