_py_collections.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543
  1. # util/_py_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. from __future__ import annotations
  9. from itertools import filterfalse
  10. from typing import AbstractSet
  11. from typing import Any
  12. from typing import Callable
  13. from typing import cast
  14. from typing import Collection
  15. from typing import Dict
  16. from typing import Iterable
  17. from typing import Iterator
  18. from typing import List
  19. from typing import Mapping
  20. from typing import NoReturn
  21. from typing import Optional
  22. from typing import Set
  23. from typing import Tuple
  24. from typing import TYPE_CHECKING
  25. from typing import TypeVar
  26. from typing import Union
  27. from ..util.typing import Self
  28. _T = TypeVar("_T", bound=Any)
  29. _S = TypeVar("_S", bound=Any)
  30. _KT = TypeVar("_KT", bound=Any)
  31. _VT = TypeVar("_VT", bound=Any)
  32. class ReadOnlyContainer:
  33. __slots__ = ()
  34. def _readonly(self, *arg: Any, **kw: Any) -> NoReturn:
  35. raise TypeError(
  36. "%s object is immutable and/or readonly" % self.__class__.__name__
  37. )
  38. def _immutable(self, *arg: Any, **kw: Any) -> NoReturn:
  39. raise TypeError("%s object is immutable" % self.__class__.__name__)
  40. def __delitem__(self, key: Any) -> NoReturn:
  41. self._readonly()
  42. def __setitem__(self, key: Any, value: Any) -> NoReturn:
  43. self._readonly()
  44. def __setattr__(self, key: str, value: Any) -> NoReturn:
  45. self._readonly()
  46. class ImmutableDictBase(ReadOnlyContainer, Dict[_KT, _VT]):
  47. if TYPE_CHECKING:
  48. def __new__(cls, *args: Any) -> Self:
  49. ...
  50. def __init__(cls, *args: Any):
  51. ...
  52. def _readonly(self, *arg: Any, **kw: Any) -> NoReturn:
  53. self._immutable()
  54. def clear(self) -> NoReturn:
  55. self._readonly()
  56. def pop(self, key: Any, default: Optional[Any] = None) -> NoReturn:
  57. self._readonly()
  58. def popitem(self) -> NoReturn:
  59. self._readonly()
  60. def setdefault(self, key: Any, default: Optional[Any] = None) -> NoReturn:
  61. self._readonly()
  62. def update(self, *arg: Any, **kw: Any) -> NoReturn:
  63. self._readonly()
  64. class immutabledict(ImmutableDictBase[_KT, _VT]):
  65. def __new__(cls, *args):
  66. new = ImmutableDictBase.__new__(cls)
  67. dict.__init__(new, *args)
  68. return new
  69. def __init__(
  70. self, *args: Union[Mapping[_KT, _VT], Iterable[Tuple[_KT, _VT]]]
  71. ):
  72. pass
  73. def __reduce__(self):
  74. return immutabledict, (dict(self),)
  75. def union(
  76. self, __d: Optional[Mapping[_KT, _VT]] = None
  77. ) -> immutabledict[_KT, _VT]:
  78. if not __d:
  79. return self
  80. new = ImmutableDictBase.__new__(self.__class__)
  81. dict.__init__(new, self)
  82. dict.update(new, __d) # type: ignore
  83. return new
  84. def _union_w_kw(
  85. self, __d: Optional[Mapping[_KT, _VT]] = None, **kw: _VT
  86. ) -> immutabledict[_KT, _VT]:
  87. # not sure if C version works correctly w/ this yet
  88. if not __d and not kw:
  89. return self
  90. new = ImmutableDictBase.__new__(self.__class__)
  91. dict.__init__(new, self)
  92. if __d:
  93. dict.update(new, __d) # type: ignore
  94. dict.update(new, kw) # type: ignore
  95. return new
  96. def merge_with(
  97. self, *dicts: Optional[Mapping[_KT, _VT]]
  98. ) -> immutabledict[_KT, _VT]:
  99. new = None
  100. for d in dicts:
  101. if d:
  102. if new is None:
  103. new = ImmutableDictBase.__new__(self.__class__)
  104. dict.__init__(new, self)
  105. dict.update(new, d) # type: ignore
  106. if new is None:
  107. return self
  108. return new
  109. def __repr__(self) -> str:
  110. return "immutabledict(%s)" % dict.__repr__(self)
  111. # PEP 584
  112. def __ior__(self, __value: Any) -> NoReturn: # type: ignore
  113. self._readonly()
  114. def __or__( # type: ignore[override]
  115. self, __value: Mapping[_KT, _VT]
  116. ) -> immutabledict[_KT, _VT]:
  117. return immutabledict(
  118. super().__or__(__value), # type: ignore[call-overload]
  119. )
  120. def __ror__( # type: ignore[override]
  121. self, __value: Mapping[_KT, _VT]
  122. ) -> immutabledict[_KT, _VT]:
  123. return immutabledict(
  124. super().__ror__(__value), # type: ignore[call-overload]
  125. )
  126. class OrderedSet(Set[_T]):
  127. __slots__ = ("_list",)
  128. _list: List[_T]
  129. def __init__(self, d: Optional[Iterable[_T]] = None) -> None:
  130. if d is not None:
  131. self._list = unique_list(d)
  132. super().update(self._list)
  133. else:
  134. self._list = []
  135. def copy(self) -> OrderedSet[_T]:
  136. cp = self.__class__()
  137. cp._list = self._list.copy()
  138. set.update(cp, cp._list)
  139. return cp
  140. def add(self, element: _T) -> None:
  141. if element not in self:
  142. self._list.append(element)
  143. super().add(element)
  144. def remove(self, element: _T) -> None:
  145. super().remove(element)
  146. self._list.remove(element)
  147. def pop(self) -> _T:
  148. try:
  149. value = self._list.pop()
  150. except IndexError:
  151. raise KeyError("pop from an empty set") from None
  152. super().remove(value)
  153. return value
  154. def insert(self, pos: int, element: _T) -> None:
  155. if element not in self:
  156. self._list.insert(pos, element)
  157. super().add(element)
  158. def discard(self, element: _T) -> None:
  159. if element in self:
  160. self._list.remove(element)
  161. super().remove(element)
  162. def clear(self) -> None:
  163. super().clear()
  164. self._list = []
  165. def __getitem__(self, key: int) -> _T:
  166. return self._list[key]
  167. def __iter__(self) -> Iterator[_T]:
  168. return iter(self._list)
  169. def __add__(self, other: Iterator[_T]) -> OrderedSet[_T]:
  170. return self.union(other)
  171. def __repr__(self) -> str:
  172. return "%s(%r)" % (self.__class__.__name__, self._list)
  173. __str__ = __repr__
  174. def update(self, *iterables: Iterable[_T]) -> None:
  175. for iterable in iterables:
  176. for e in iterable:
  177. if e not in self:
  178. self._list.append(e)
  179. super().add(e)
  180. def __ior__(self, other: AbstractSet[_S]) -> OrderedSet[Union[_T, _S]]:
  181. self.update(other)
  182. return self
  183. def union(self, *other: Iterable[_S]) -> OrderedSet[Union[_T, _S]]:
  184. result: OrderedSet[Union[_T, _S]] = self.copy()
  185. result.update(*other)
  186. return result
  187. def __or__(self, other: AbstractSet[_S]) -> OrderedSet[Union[_T, _S]]:
  188. return self.union(other)
  189. def intersection(self, *other: Iterable[Any]) -> OrderedSet[_T]:
  190. other_set: Set[Any] = set()
  191. other_set.update(*other)
  192. return self.__class__(a for a in self if a in other_set)
  193. def __and__(self, other: AbstractSet[object]) -> OrderedSet[_T]:
  194. return self.intersection(other)
  195. def symmetric_difference(self, other: Iterable[_T]) -> OrderedSet[_T]:
  196. collection: Collection[_T]
  197. if isinstance(other, set):
  198. collection = other_set = other
  199. elif isinstance(other, Collection):
  200. collection = other
  201. other_set = set(other)
  202. else:
  203. collection = list(other)
  204. other_set = set(collection)
  205. result = self.__class__(a for a in self if a not in other_set)
  206. result.update(a for a in collection if a not in self)
  207. return result
  208. def __xor__(self, other: AbstractSet[_S]) -> OrderedSet[Union[_T, _S]]:
  209. return cast(OrderedSet[Union[_T, _S]], self).symmetric_difference(
  210. other
  211. )
  212. def difference(self, *other: Iterable[Any]) -> OrderedSet[_T]:
  213. other_set = super().difference(*other)
  214. return self.__class__(a for a in self._list if a in other_set)
  215. def __sub__(self, other: AbstractSet[Optional[_T]]) -> OrderedSet[_T]:
  216. return self.difference(other)
  217. def intersection_update(self, *other: Iterable[Any]) -> None:
  218. super().intersection_update(*other)
  219. self._list = [a for a in self._list if a in self]
  220. def __iand__(self, other: AbstractSet[object]) -> OrderedSet[_T]:
  221. self.intersection_update(other)
  222. return self
  223. def symmetric_difference_update(self, other: Iterable[Any]) -> None:
  224. collection = other if isinstance(other, Collection) else list(other)
  225. super().symmetric_difference_update(collection)
  226. self._list = [a for a in self._list if a in self]
  227. self._list += [a for a in collection if a in self]
  228. def __ixor__(self, other: AbstractSet[_S]) -> OrderedSet[Union[_T, _S]]:
  229. self.symmetric_difference_update(other)
  230. return cast(OrderedSet[Union[_T, _S]], self)
  231. def difference_update(self, *other: Iterable[Any]) -> None:
  232. super().difference_update(*other)
  233. self._list = [a for a in self._list if a in self]
  234. def __isub__(self, other: AbstractSet[Optional[_T]]) -> OrderedSet[_T]: # type: ignore # noqa: E501
  235. self.difference_update(other)
  236. return self
  237. class IdentitySet:
  238. """A set that considers only object id() for uniqueness.
  239. This strategy has edge cases for builtin types- it's possible to have
  240. two 'foo' strings in one of these sets, for example. Use sparingly.
  241. """
  242. _members: Dict[int, Any]
  243. def __init__(self, iterable: Optional[Iterable[Any]] = None):
  244. self._members = dict()
  245. if iterable:
  246. self.update(iterable)
  247. def add(self, value: Any) -> None:
  248. self._members[id(value)] = value
  249. def __contains__(self, value: Any) -> bool:
  250. return id(value) in self._members
  251. def remove(self, value: Any) -> None:
  252. del self._members[id(value)]
  253. def discard(self, value: Any) -> None:
  254. try:
  255. self.remove(value)
  256. except KeyError:
  257. pass
  258. def pop(self) -> Any:
  259. try:
  260. pair = self._members.popitem()
  261. return pair[1]
  262. except KeyError:
  263. raise KeyError("pop from an empty set")
  264. def clear(self) -> None:
  265. self._members.clear()
  266. def __eq__(self, other: Any) -> bool:
  267. if isinstance(other, IdentitySet):
  268. return self._members == other._members
  269. else:
  270. return False
  271. def __ne__(self, other: Any) -> bool:
  272. if isinstance(other, IdentitySet):
  273. return self._members != other._members
  274. else:
  275. return True
  276. def issubset(self, iterable: Iterable[Any]) -> bool:
  277. if isinstance(iterable, self.__class__):
  278. other = iterable
  279. else:
  280. other = self.__class__(iterable)
  281. if len(self) > len(other):
  282. return False
  283. for m in filterfalse(
  284. other._members.__contains__, iter(self._members.keys())
  285. ):
  286. return False
  287. return True
  288. def __le__(self, other: Any) -> bool:
  289. if not isinstance(other, IdentitySet):
  290. return NotImplemented
  291. return self.issubset(other)
  292. def __lt__(self, other: Any) -> bool:
  293. if not isinstance(other, IdentitySet):
  294. return NotImplemented
  295. return len(self) < len(other) and self.issubset(other)
  296. def issuperset(self, iterable: Iterable[Any]) -> bool:
  297. if isinstance(iterable, self.__class__):
  298. other = iterable
  299. else:
  300. other = self.__class__(iterable)
  301. if len(self) < len(other):
  302. return False
  303. for m in filterfalse(
  304. self._members.__contains__, iter(other._members.keys())
  305. ):
  306. return False
  307. return True
  308. def __ge__(self, other: Any) -> bool:
  309. if not isinstance(other, IdentitySet):
  310. return NotImplemented
  311. return self.issuperset(other)
  312. def __gt__(self, other: Any) -> bool:
  313. if not isinstance(other, IdentitySet):
  314. return NotImplemented
  315. return len(self) > len(other) and self.issuperset(other)
  316. def union(self, iterable: Iterable[Any]) -> IdentitySet:
  317. result = self.__class__()
  318. members = self._members
  319. result._members.update(members)
  320. result._members.update((id(obj), obj) for obj in iterable)
  321. return result
  322. def __or__(self, other: Any) -> IdentitySet:
  323. if not isinstance(other, IdentitySet):
  324. return NotImplemented
  325. return self.union(other)
  326. def update(self, iterable: Iterable[Any]) -> None:
  327. self._members.update((id(obj), obj) for obj in iterable)
  328. def __ior__(self, other: Any) -> IdentitySet:
  329. if not isinstance(other, IdentitySet):
  330. return NotImplemented
  331. self.update(other)
  332. return self
  333. def difference(self, iterable: Iterable[Any]) -> IdentitySet:
  334. result = self.__new__(self.__class__)
  335. other: Collection[Any]
  336. if isinstance(iterable, self.__class__):
  337. other = iterable._members
  338. else:
  339. other = {id(obj) for obj in iterable}
  340. result._members = {
  341. k: v for k, v in self._members.items() if k not in other
  342. }
  343. return result
  344. def __sub__(self, other: IdentitySet) -> IdentitySet:
  345. if not isinstance(other, IdentitySet):
  346. return NotImplemented
  347. return self.difference(other)
  348. def difference_update(self, iterable: Iterable[Any]) -> None:
  349. self._members = self.difference(iterable)._members
  350. def __isub__(self, other: IdentitySet) -> IdentitySet:
  351. if not isinstance(other, IdentitySet):
  352. return NotImplemented
  353. self.difference_update(other)
  354. return self
  355. def intersection(self, iterable: Iterable[Any]) -> IdentitySet:
  356. result = self.__new__(self.__class__)
  357. other: Collection[Any]
  358. if isinstance(iterable, self.__class__):
  359. other = iterable._members
  360. else:
  361. other = {id(obj) for obj in iterable}
  362. result._members = {
  363. k: v for k, v in self._members.items() if k in other
  364. }
  365. return result
  366. def __and__(self, other: IdentitySet) -> IdentitySet:
  367. if not isinstance(other, IdentitySet):
  368. return NotImplemented
  369. return self.intersection(other)
  370. def intersection_update(self, iterable: Iterable[Any]) -> None:
  371. self._members = self.intersection(iterable)._members
  372. def __iand__(self, other: IdentitySet) -> IdentitySet:
  373. if not isinstance(other, IdentitySet):
  374. return NotImplemented
  375. self.intersection_update(other)
  376. return self
  377. def symmetric_difference(self, iterable: Iterable[Any]) -> IdentitySet:
  378. result = self.__new__(self.__class__)
  379. if isinstance(iterable, self.__class__):
  380. other = iterable._members
  381. else:
  382. other = {id(obj): obj for obj in iterable}
  383. result._members = {
  384. k: v for k, v in self._members.items() if k not in other
  385. }
  386. result._members.update(
  387. (k, v) for k, v in other.items() if k not in self._members
  388. )
  389. return result
  390. def __xor__(self, other: IdentitySet) -> IdentitySet:
  391. if not isinstance(other, IdentitySet):
  392. return NotImplemented
  393. return self.symmetric_difference(other)
  394. def symmetric_difference_update(self, iterable: Iterable[Any]) -> None:
  395. self._members = self.symmetric_difference(iterable)._members
  396. def __ixor__(self, other: IdentitySet) -> IdentitySet:
  397. if not isinstance(other, IdentitySet):
  398. return NotImplemented
  399. self.symmetric_difference(other)
  400. return self
  401. def copy(self) -> IdentitySet:
  402. result = self.__new__(self.__class__)
  403. result._members = self._members.copy()
  404. return result
  405. __copy__ = copy
  406. def __len__(self) -> int:
  407. return len(self._members)
  408. def __iter__(self) -> Iterator[Any]:
  409. return iter(self._members.values())
  410. def __hash__(self) -> NoReturn:
  411. raise TypeError("set objects are unhashable")
  412. def __repr__(self) -> str:
  413. return "%s(%r)" % (type(self).__name__, list(self._members.values()))
  414. def unique_list(
  415. seq: Iterable[_T], hashfunc: Optional[Callable[[_T], int]] = None
  416. ) -> List[_T]:
  417. seen: Set[Any] = set()
  418. seen_add = seen.add
  419. if not hashfunc:
  420. return [x for x in seq if x not in seen and not seen_add(x)]
  421. else:
  422. return [
  423. x
  424. for x in seq
  425. if hashfunc(x) not in seen and not seen_add(hashfunc(x))
  426. ]