_collections.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721
  1. # util/_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
  8. """Collection classes and helpers."""
  9. from __future__ import annotations
  10. import operator
  11. import threading
  12. import types
  13. import typing
  14. from typing import Any
  15. from typing import Callable
  16. from typing import cast
  17. from typing import Dict
  18. from typing import FrozenSet
  19. from typing import Generic
  20. from typing import Iterable
  21. from typing import Iterator
  22. from typing import List
  23. from typing import Mapping
  24. from typing import NoReturn
  25. from typing import Optional
  26. from typing import overload
  27. from typing import Sequence
  28. from typing import Set
  29. from typing import Tuple
  30. from typing import TypeVar
  31. from typing import Union
  32. from typing import ValuesView
  33. import weakref
  34. from ._has_cy import HAS_CYEXTENSION
  35. from .typing import is_non_string_iterable
  36. from .typing import Literal
  37. from .typing import Protocol
  38. if typing.TYPE_CHECKING or not HAS_CYEXTENSION:
  39. from ._py_collections import immutabledict as immutabledict
  40. from ._py_collections import IdentitySet as IdentitySet
  41. from ._py_collections import ReadOnlyContainer as ReadOnlyContainer
  42. from ._py_collections import ImmutableDictBase as ImmutableDictBase
  43. from ._py_collections import OrderedSet as OrderedSet
  44. from ._py_collections import unique_list as unique_list
  45. else:
  46. from sqlalchemy.cyextension.immutabledict import (
  47. ReadOnlyContainer as ReadOnlyContainer,
  48. )
  49. from sqlalchemy.cyextension.immutabledict import (
  50. ImmutableDictBase as ImmutableDictBase,
  51. )
  52. from sqlalchemy.cyextension.immutabledict import (
  53. immutabledict as immutabledict,
  54. )
  55. from sqlalchemy.cyextension.collections import IdentitySet as IdentitySet
  56. from sqlalchemy.cyextension.collections import OrderedSet as OrderedSet
  57. from sqlalchemy.cyextension.collections import ( # noqa
  58. unique_list as unique_list,
  59. )
  60. _T = TypeVar("_T", bound=Any)
  61. _KT = TypeVar("_KT", bound=Any)
  62. _VT = TypeVar("_VT", bound=Any)
  63. _T_co = TypeVar("_T_co", covariant=True)
  64. EMPTY_SET: FrozenSet[Any] = frozenset()
  65. NONE_SET: FrozenSet[Any] = frozenset([None])
  66. def merge_lists_w_ordering(a: List[Any], b: List[Any]) -> List[Any]:
  67. """merge two lists, maintaining ordering as much as possible.
  68. this is to reconcile vars(cls) with cls.__annotations__.
  69. Example::
  70. >>> a = ['__tablename__', 'id', 'x', 'created_at']
  71. >>> b = ['id', 'name', 'data', 'y', 'created_at']
  72. >>> merge_lists_w_ordering(a, b)
  73. ['__tablename__', 'id', 'name', 'data', 'y', 'x', 'created_at']
  74. This is not necessarily the ordering that things had on the class,
  75. in this case the class is::
  76. class User(Base):
  77. __tablename__ = "users"
  78. id: Mapped[int] = mapped_column(primary_key=True)
  79. name: Mapped[str]
  80. data: Mapped[Optional[str]]
  81. x = Column(Integer)
  82. y: Mapped[int]
  83. created_at: Mapped[datetime.datetime] = mapped_column()
  84. But things are *mostly* ordered.
  85. The algorithm could also be done by creating a partial ordering for
  86. all items in both lists and then using topological_sort(), but that
  87. is too much overhead.
  88. Background on how I came up with this is at:
  89. https://gist.github.com/zzzeek/89de958cf0803d148e74861bd682ebae
  90. """
  91. overlap = set(a).intersection(b)
  92. result = []
  93. current, other = iter(a), iter(b)
  94. while True:
  95. for element in current:
  96. if element in overlap:
  97. overlap.discard(element)
  98. other, current = current, other
  99. break
  100. result.append(element)
  101. else:
  102. result.extend(other)
  103. break
  104. return result
  105. def coerce_to_immutabledict(d: Mapping[_KT, _VT]) -> immutabledict[_KT, _VT]:
  106. if not d:
  107. return EMPTY_DICT
  108. elif isinstance(d, immutabledict):
  109. return d
  110. else:
  111. return immutabledict(d)
  112. EMPTY_DICT: immutabledict[Any, Any] = immutabledict()
  113. class FacadeDict(ImmutableDictBase[_KT, _VT]):
  114. """A dictionary that is not publicly mutable."""
  115. def __new__(cls, *args: Any) -> FacadeDict[Any, Any]:
  116. new = ImmutableDictBase.__new__(cls)
  117. return new
  118. def copy(self) -> NoReturn:
  119. raise NotImplementedError(
  120. "an immutabledict shouldn't need to be copied. use dict(d) "
  121. "if you need a mutable dictionary."
  122. )
  123. def __reduce__(self) -> Any:
  124. return FacadeDict, (dict(self),)
  125. def _insert_item(self, key: _KT, value: _VT) -> None:
  126. """insert an item into the dictionary directly."""
  127. dict.__setitem__(self, key, value)
  128. def __repr__(self) -> str:
  129. return "FacadeDict(%s)" % dict.__repr__(self)
  130. _DT = TypeVar("_DT", bound=Any)
  131. _F = TypeVar("_F", bound=Any)
  132. class Properties(Generic[_T]):
  133. """Provide a __getattr__/__setattr__ interface over a dict."""
  134. __slots__ = ("_data",)
  135. _data: Dict[str, _T]
  136. def __init__(self, data: Dict[str, _T]):
  137. object.__setattr__(self, "_data", data)
  138. def __len__(self) -> int:
  139. return len(self._data)
  140. def __iter__(self) -> Iterator[_T]:
  141. return iter(list(self._data.values()))
  142. def __dir__(self) -> List[str]:
  143. return dir(super()) + [str(k) for k in self._data.keys()]
  144. def __add__(self, other: Properties[_F]) -> List[Union[_T, _F]]:
  145. return list(self) + list(other)
  146. def __setitem__(self, key: str, obj: _T) -> None:
  147. self._data[key] = obj
  148. def __getitem__(self, key: str) -> _T:
  149. return self._data[key]
  150. def __delitem__(self, key: str) -> None:
  151. del self._data[key]
  152. def __setattr__(self, key: str, obj: _T) -> None:
  153. self._data[key] = obj
  154. def __getstate__(self) -> Dict[str, Any]:
  155. return {"_data": self._data}
  156. def __setstate__(self, state: Dict[str, Any]) -> None:
  157. object.__setattr__(self, "_data", state["_data"])
  158. def __getattr__(self, key: str) -> _T:
  159. try:
  160. return self._data[key]
  161. except KeyError:
  162. raise AttributeError(key)
  163. def __contains__(self, key: str) -> bool:
  164. return key in self._data
  165. def as_readonly(self) -> ReadOnlyProperties[_T]:
  166. """Return an immutable proxy for this :class:`.Properties`."""
  167. return ReadOnlyProperties(self._data)
  168. def update(self, value: Dict[str, _T]) -> None:
  169. self._data.update(value)
  170. @overload
  171. def get(self, key: str) -> Optional[_T]:
  172. ...
  173. @overload
  174. def get(self, key: str, default: Union[_DT, _T]) -> Union[_DT, _T]:
  175. ...
  176. def get(
  177. self, key: str, default: Optional[Union[_DT, _T]] = None
  178. ) -> Optional[Union[_T, _DT]]:
  179. if key in self:
  180. return self[key]
  181. else:
  182. return default
  183. def keys(self) -> List[str]:
  184. return list(self._data)
  185. def values(self) -> List[_T]:
  186. return list(self._data.values())
  187. def items(self) -> List[Tuple[str, _T]]:
  188. return list(self._data.items())
  189. def has_key(self, key: str) -> bool:
  190. return key in self._data
  191. def clear(self) -> None:
  192. self._data.clear()
  193. class OrderedProperties(Properties[_T]):
  194. """Provide a __getattr__/__setattr__ interface with an OrderedDict
  195. as backing store."""
  196. __slots__ = ()
  197. def __init__(self):
  198. Properties.__init__(self, OrderedDict())
  199. class ReadOnlyProperties(ReadOnlyContainer, Properties[_T]):
  200. """Provide immutable dict/object attribute to an underlying dictionary."""
  201. __slots__ = ()
  202. def _ordered_dictionary_sort(d, key=None):
  203. """Sort an OrderedDict in-place."""
  204. items = [(k, d[k]) for k in sorted(d, key=key)]
  205. d.clear()
  206. d.update(items)
  207. OrderedDict = dict
  208. sort_dictionary = _ordered_dictionary_sort
  209. class WeakSequence(Sequence[_T]):
  210. def __init__(self, __elements: Sequence[_T] = ()):
  211. # adapted from weakref.WeakKeyDictionary, prevent reference
  212. # cycles in the collection itself
  213. def _remove(item, selfref=weakref.ref(self)):
  214. self = selfref()
  215. if self is not None:
  216. self._storage.remove(item)
  217. self._remove = _remove
  218. self._storage = [
  219. weakref.ref(element, _remove) for element in __elements
  220. ]
  221. def append(self, item):
  222. self._storage.append(weakref.ref(item, self._remove))
  223. def __len__(self):
  224. return len(self._storage)
  225. def __iter__(self):
  226. return (
  227. obj for obj in (ref() for ref in self._storage) if obj is not None
  228. )
  229. def __getitem__(self, index):
  230. try:
  231. obj = self._storage[index]
  232. except KeyError:
  233. raise IndexError("Index %s out of range" % index)
  234. else:
  235. return obj()
  236. class OrderedIdentitySet(IdentitySet):
  237. def __init__(self, iterable: Optional[Iterable[Any]] = None):
  238. IdentitySet.__init__(self)
  239. self._members = OrderedDict()
  240. if iterable:
  241. for o in iterable:
  242. self.add(o)
  243. class PopulateDict(Dict[_KT, _VT]):
  244. """A dict which populates missing values via a creation function.
  245. Note the creation function takes a key, unlike
  246. collections.defaultdict.
  247. """
  248. def __init__(self, creator: Callable[[_KT], _VT]):
  249. self.creator = creator
  250. def __missing__(self, key: Any) -> Any:
  251. self[key] = val = self.creator(key)
  252. return val
  253. class WeakPopulateDict(Dict[_KT, _VT]):
  254. """Like PopulateDict, but assumes a self + a method and does not create
  255. a reference cycle.
  256. """
  257. def __init__(self, creator_method: types.MethodType):
  258. self.creator = creator_method.__func__
  259. weakself = creator_method.__self__
  260. self.weakself = weakref.ref(weakself)
  261. def __missing__(self, key: Any) -> Any:
  262. self[key] = val = self.creator(self.weakself(), key)
  263. return val
  264. # Define collections that are capable of storing
  265. # ColumnElement objects as hashable keys/elements.
  266. # At this point, these are mostly historical, things
  267. # used to be more complicated.
  268. column_set = set
  269. column_dict = dict
  270. ordered_column_set = OrderedSet
  271. class UniqueAppender(Generic[_T]):
  272. """Appends items to a collection ensuring uniqueness.
  273. Additional appends() of the same object are ignored. Membership is
  274. determined by identity (``is a``) not equality (``==``).
  275. """
  276. __slots__ = "data", "_data_appender", "_unique"
  277. data: Union[Iterable[_T], Set[_T], List[_T]]
  278. _data_appender: Callable[[_T], None]
  279. _unique: Dict[int, Literal[True]]
  280. def __init__(
  281. self,
  282. data: Union[Iterable[_T], Set[_T], List[_T]],
  283. via: Optional[str] = None,
  284. ):
  285. self.data = data
  286. self._unique = {}
  287. if via:
  288. self._data_appender = getattr(data, via)
  289. elif hasattr(data, "append"):
  290. self._data_appender = cast("List[_T]", data).append
  291. elif hasattr(data, "add"):
  292. self._data_appender = cast("Set[_T]", data).add
  293. def append(self, item: _T) -> None:
  294. id_ = id(item)
  295. if id_ not in self._unique:
  296. self._data_appender(item)
  297. self._unique[id_] = True
  298. def __iter__(self) -> Iterator[_T]:
  299. return iter(self.data)
  300. def coerce_generator_arg(arg: Any) -> List[Any]:
  301. if len(arg) == 1 and isinstance(arg[0], types.GeneratorType):
  302. return list(arg[0])
  303. else:
  304. return cast("List[Any]", arg)
  305. def to_list(x: Any, default: Optional[List[Any]] = None) -> List[Any]:
  306. if x is None:
  307. return default # type: ignore
  308. if not is_non_string_iterable(x):
  309. return [x]
  310. elif isinstance(x, list):
  311. return x
  312. else:
  313. return list(x)
  314. def has_intersection(set_, iterable):
  315. r"""return True if any items of set\_ are present in iterable.
  316. Goes through special effort to ensure __hash__ is not called
  317. on items in iterable that don't support it.
  318. """
  319. # TODO: optimize, write in C, etc.
  320. return bool(set_.intersection([i for i in iterable if i.__hash__]))
  321. def to_set(x):
  322. if x is None:
  323. return set()
  324. if not isinstance(x, set):
  325. return set(to_list(x))
  326. else:
  327. return x
  328. def to_column_set(x: Any) -> Set[Any]:
  329. if x is None:
  330. return column_set()
  331. if not isinstance(x, column_set):
  332. return column_set(to_list(x))
  333. else:
  334. return x
  335. def update_copy(d, _new=None, **kw):
  336. """Copy the given dict and update with the given values."""
  337. d = d.copy()
  338. if _new:
  339. d.update(_new)
  340. d.update(**kw)
  341. return d
  342. def flatten_iterator(x: Iterable[_T]) -> Iterator[_T]:
  343. """Given an iterator of which further sub-elements may also be
  344. iterators, flatten the sub-elements into a single iterator.
  345. """
  346. elem: _T
  347. for elem in x:
  348. if not isinstance(elem, str) and hasattr(elem, "__iter__"):
  349. yield from flatten_iterator(elem)
  350. else:
  351. yield elem
  352. class LRUCache(typing.MutableMapping[_KT, _VT]):
  353. """Dictionary with 'squishy' removal of least
  354. recently used items.
  355. Note that either get() or [] should be used here, but
  356. generally its not safe to do an "in" check first as the dictionary
  357. can change subsequent to that call.
  358. """
  359. __slots__ = (
  360. "capacity",
  361. "threshold",
  362. "size_alert",
  363. "_data",
  364. "_counter",
  365. "_mutex",
  366. )
  367. capacity: int
  368. threshold: float
  369. size_alert: Optional[Callable[[LRUCache[_KT, _VT]], None]]
  370. def __init__(
  371. self,
  372. capacity: int = 100,
  373. threshold: float = 0.5,
  374. size_alert: Optional[Callable[..., None]] = None,
  375. ):
  376. self.capacity = capacity
  377. self.threshold = threshold
  378. self.size_alert = size_alert
  379. self._counter = 0
  380. self._mutex = threading.Lock()
  381. self._data: Dict[_KT, Tuple[_KT, _VT, List[int]]] = {}
  382. def _inc_counter(self):
  383. self._counter += 1
  384. return self._counter
  385. @overload
  386. def get(self, key: _KT) -> Optional[_VT]:
  387. ...
  388. @overload
  389. def get(self, key: _KT, default: Union[_VT, _T]) -> Union[_VT, _T]:
  390. ...
  391. def get(
  392. self, key: _KT, default: Optional[Union[_VT, _T]] = None
  393. ) -> Optional[Union[_VT, _T]]:
  394. item = self._data.get(key)
  395. if item is not None:
  396. item[2][0] = self._inc_counter()
  397. return item[1]
  398. else:
  399. return default
  400. def __getitem__(self, key: _KT) -> _VT:
  401. item = self._data[key]
  402. item[2][0] = self._inc_counter()
  403. return item[1]
  404. def __iter__(self) -> Iterator[_KT]:
  405. return iter(self._data)
  406. def __len__(self) -> int:
  407. return len(self._data)
  408. def values(self) -> ValuesView[_VT]:
  409. return typing.ValuesView({k: i[1] for k, i in self._data.items()})
  410. def __setitem__(self, key: _KT, value: _VT) -> None:
  411. self._data[key] = (key, value, [self._inc_counter()])
  412. self._manage_size()
  413. def __delitem__(self, __v: _KT) -> None:
  414. del self._data[__v]
  415. @property
  416. def size_threshold(self) -> float:
  417. return self.capacity + self.capacity * self.threshold
  418. def _manage_size(self) -> None:
  419. if not self._mutex.acquire(False):
  420. return
  421. try:
  422. size_alert = bool(self.size_alert)
  423. while len(self) > self.capacity + self.capacity * self.threshold:
  424. if size_alert:
  425. size_alert = False
  426. self.size_alert(self) # type: ignore
  427. by_counter = sorted(
  428. self._data.values(),
  429. key=operator.itemgetter(2),
  430. reverse=True,
  431. )
  432. for item in by_counter[self.capacity :]:
  433. try:
  434. del self._data[item[0]]
  435. except KeyError:
  436. # deleted elsewhere; skip
  437. continue
  438. finally:
  439. self._mutex.release()
  440. class _CreateFuncType(Protocol[_T_co]):
  441. def __call__(self) -> _T_co:
  442. ...
  443. class _ScopeFuncType(Protocol):
  444. def __call__(self) -> Any:
  445. ...
  446. class ScopedRegistry(Generic[_T]):
  447. """A Registry that can store one or multiple instances of a single
  448. class on the basis of a "scope" function.
  449. The object implements ``__call__`` as the "getter", so by
  450. calling ``myregistry()`` the contained object is returned
  451. for the current scope.
  452. :param createfunc:
  453. a callable that returns a new object to be placed in the registry
  454. :param scopefunc:
  455. a callable that will return a key to store/retrieve an object.
  456. """
  457. __slots__ = "createfunc", "scopefunc", "registry"
  458. createfunc: _CreateFuncType[_T]
  459. scopefunc: _ScopeFuncType
  460. registry: Any
  461. def __init__(
  462. self, createfunc: Callable[[], _T], scopefunc: Callable[[], Any]
  463. ):
  464. """Construct a new :class:`.ScopedRegistry`.
  465. :param createfunc: A creation function that will generate
  466. a new value for the current scope, if none is present.
  467. :param scopefunc: A function that returns a hashable
  468. token representing the current scope (such as, current
  469. thread identifier).
  470. """
  471. self.createfunc = createfunc
  472. self.scopefunc = scopefunc
  473. self.registry = {}
  474. def __call__(self) -> _T:
  475. key = self.scopefunc()
  476. try:
  477. return self.registry[key] # type: ignore[no-any-return]
  478. except KeyError:
  479. return self.registry.setdefault(key, self.createfunc()) # type: ignore[no-any-return] # noqa: E501
  480. def has(self) -> bool:
  481. """Return True if an object is present in the current scope."""
  482. return self.scopefunc() in self.registry
  483. def set(self, obj: _T) -> None:
  484. """Set the value for the current scope."""
  485. self.registry[self.scopefunc()] = obj
  486. def clear(self) -> None:
  487. """Clear the current scope, if any."""
  488. try:
  489. del self.registry[self.scopefunc()]
  490. except KeyError:
  491. pass
  492. class ThreadLocalRegistry(ScopedRegistry[_T]):
  493. """A :class:`.ScopedRegistry` that uses a ``threading.local()``
  494. variable for storage.
  495. """
  496. def __init__(self, createfunc: Callable[[], _T]):
  497. self.createfunc = createfunc
  498. self.registry = threading.local()
  499. def __call__(self) -> _T:
  500. try:
  501. return self.registry.value # type: ignore[no-any-return]
  502. except AttributeError:
  503. val = self.registry.value = self.createfunc()
  504. return val
  505. def has(self) -> bool:
  506. return hasattr(self.registry, "value")
  507. def set(self, obj: _T) -> None:
  508. self.registry.value = obj
  509. def clear(self) -> None:
  510. try:
  511. del self.registry.value
  512. except AttributeError:
  513. pass
  514. def has_dupes(sequence, target):
  515. """Given a sequence and search object, return True if there's more
  516. than one, False if zero or one of them.
  517. """
  518. # compare to .index version below, this version introduces less function
  519. # overhead and is usually the same speed. At 15000 items (way bigger than
  520. # a relationship-bound collection in memory usually is) it begins to
  521. # fall behind the other version only by microseconds.
  522. c = 0
  523. for item in sequence:
  524. if item is target:
  525. c += 1
  526. if c > 1:
  527. return True
  528. return False
  529. # .index version. the two __contains__ calls as well
  530. # as .index() and isinstance() slow this down.
  531. # def has_dupes(sequence, target):
  532. # if target not in sequence:
  533. # return False
  534. # elif not isinstance(sequence, collections_abc.Sequence):
  535. # return False
  536. #
  537. # idx = sequence.index(target)
  538. # return target in sequence[idx + 1:]