Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# -*- coding: utf-8 -*- 

2"""Custom maps, sets, sequences, and other data structures.""" 

3from __future__ import absolute_import, unicode_literals 

4 

5import sys 

6from collections import OrderedDict as _OrderedDict 

7from collections import deque 

8from heapq import heapify, heappop, heappush 

9from itertools import chain, count 

10 

11from celery.five import (PY3, Empty, items, keys, monotonic, 

12 python_2_unicode_compatible, values) 

13 

14from .functional import first, uniq 

15from .text import match_case 

16 

17try: 

18 from collections.abc import Callable, Mapping, MutableMapping, MutableSet 

19 from collections.abc import Sequence 

20except ImportError: 

21 # TODO: Remove this when we drop Python 2.7 support 

22 from collections import Callable, Mapping, MutableMapping, MutableSet 

23 from collections import Sequence 

24 

25 

26try: 

27 # pypy: dicts are ordered in recent versions 

28 from __pypy__ import reversed_dict as _dict_is_ordered 

29except ImportError: 

30 _dict_is_ordered = None 

31 

32try: 

33 from django.utils.functional import LazyObject, LazySettings 

34except ImportError: 

35 class LazyObject(object): # noqa 

36 pass 

37 LazySettings = LazyObject # noqa 

38 

39__all__ = ( 

40 'AttributeDictMixin', 'AttributeDict', 'BufferMap', 'ChainMap', 

41 'ConfigurationView', 'DictAttribute', 'Evictable', 

42 'LimitedSet', 'Messagebuffer', 'OrderedDict', 

43 'force_mapping', 'lpmerge', 

44) 

45 

46REPR_LIMITED_SET = """\ 

47<{name}({size}): maxlen={0.maxlen}, expires={0.expires}, minlen={0.minlen}>\ 

48""" 

49 

50 

51def force_mapping(m): 

52 # type: (Any) -> Mapping 

53 """Wrap object into supporting the mapping interface if necessary.""" 

54 if isinstance(m, (LazyObject, LazySettings)): 

55 m = m._wrapped 

56 return DictAttribute(m) if not isinstance(m, Mapping) else m 

57 

58 

59def lpmerge(L, R): 

60 # type: (Mapping, Mapping) -> Mapping 

61 """In place left precedent dictionary merge. 

62 

63 Keeps values from `L`, if the value in `R` is :const:`None`. 

64 """ 

65 setitem = L.__setitem__ 

66 [setitem(k, v) for k, v in items(R) if v is not None] 

67 return L 

68 

69 

70class OrderedDict(_OrderedDict): 

71 """Dict where insertion order matters.""" 

72 

73 if PY3: # pragma: no cover 

74 def _LRUkey(self): 

75 # type: () -> Any 

76 # return value of od.keys does not support __next__, 

77 # but this version will also not create a copy of the list. 

78 return next(iter(keys(self))) 

79 else: 

80 if _dict_is_ordered: # pragma: no cover 

81 def _LRUkey(self): 

82 # type: () -> Any 

83 # iterkeys is iterable. 

84 return next(self.iterkeys()) 

85 else: 

86 def _LRUkey(self): 

87 # type: () -> Any 

88 return self._OrderedDict__root[1][2] 

89 

90 if not hasattr(_OrderedDict, 'move_to_end'): 

91 if _dict_is_ordered: # pragma: no cover 

92 

93 def move_to_end(self, key, last=True): 

94 # type: (Any, bool) -> None 

95 if not last: 

96 # we don't use this argument, and the only way to 

97 # implement this on PyPy seems to be O(n): creating a 

98 # copy with the order changed, so we just raise. 

99 raise NotImplementedError('no last=True on PyPy') 

100 self[key] = self.pop(key) 

101 

102 else: 

103 

104 def move_to_end(self, key, last=True): 

105 # type: (Any, bool) -> None 

106 link = self._OrderedDict__map[key] 

107 link_prev = link[0] 

108 link_next = link[1] 

109 link_prev[1] = link_next 

110 link_next[0] = link_prev 

111 root = self._OrderedDict__root 

112 if last: 

113 last = root[0] 

114 link[0] = last 

115 link[1] = root 

116 last[1] = root[0] = link 

117 else: 

118 first_node = root[1] 

119 link[0] = root 

120 link[1] = first_node 

121 root[1] = first_node[0] = link 

122 

123 

124class AttributeDictMixin(object): 

125 """Mixin for Mapping interface that adds attribute access. 

126 

127 I.e., `d.key -> d[key]`). 

128 """ 

129 

130 def __getattr__(self, k): 

131 # type: (str) -> Any 

132 """`d.key -> d[key]`.""" 

133 try: 

134 return self[k] 

135 except KeyError: 

136 raise AttributeError( 

137 '{0!r} object has no attribute {1!r}'.format( 

138 type(self).__name__, k)) 

139 

140 def __setattr__(self, key, value): 

141 # type: (str, Any) -> None 

142 """`d[key] = value -> d.key = value`.""" 

143 self[key] = value 

144 

145 

146class AttributeDict(dict, AttributeDictMixin): 

147 """Dict subclass with attribute access.""" 

148 

149 

150class DictAttribute(object): 

151 """Dict interface to attributes. 

152 

153 `obj[k] -> obj.k` 

154 `obj[k] = val -> obj.k = val` 

155 """ 

156 

157 obj = None 

158 

159 def __init__(self, obj): 

160 # type: (Any) -> None 

161 object.__setattr__(self, 'obj', obj) 

162 

163 def __getattr__(self, key): 

164 # type: (Any) -> Any 

165 return getattr(self.obj, key) 

166 

167 def __setattr__(self, key, value): 

168 # type: (Any, Any) -> None 

169 return setattr(self.obj, key, value) 

170 

171 def get(self, key, default=None): 

172 # type: (Any, Any) -> Any 

173 try: 

174 return self[key] 

175 except KeyError: 

176 return default 

177 

178 def setdefault(self, key, default=None): 

179 # type: (Any, Any) -> None 

180 if key not in self: 

181 self[key] = default 

182 

183 def __getitem__(self, key): 

184 # type: (Any) -> Any 

185 try: 

186 return getattr(self.obj, key) 

187 except AttributeError: 

188 raise KeyError(key) 

189 

190 def __setitem__(self, key, value): 

191 # type: (Any, Any) -> Any 

192 setattr(self.obj, key, value) 

193 

194 def __contains__(self, key): 

195 # type: (Any) -> bool 

196 return hasattr(self.obj, key) 

197 

198 def _iterate_keys(self): 

199 # type: () -> Iterable 

200 return iter(dir(self.obj)) 

201 iterkeys = _iterate_keys 

202 

203 def __iter__(self): 

204 # type: () -> Iterable 

205 return self._iterate_keys() 

206 

207 def _iterate_items(self): 

208 # type: () -> Iterable 

209 for key in self._iterate_keys(): 

210 yield key, getattr(self.obj, key) 

211 iteritems = _iterate_items 

212 

213 def _iterate_values(self): 

214 # type: () -> Iterable 

215 for key in self._iterate_keys(): 

216 yield getattr(self.obj, key) 

217 itervalues = _iterate_values 

218 

219 if sys.version_info[0] == 3: # pragma: no cover 

220 items = _iterate_items 

221 keys = _iterate_keys 

222 values = _iterate_values 

223 else: 

224 

225 def keys(self): 

226 # type: () -> List[Any] 

227 return list(self) 

228 

229 def items(self): 

230 # type: () -> List[Tuple[Any, Any]] 

231 return list(self._iterate_items()) 

232 

233 def values(self): 

234 # type: () -> List[Any] 

235 return list(self._iterate_values()) 

236 

237 

238MutableMapping.register(DictAttribute) # noqa: E305 

239 

240 

241class ChainMap(MutableMapping): 

242 """Key lookup on a sequence of maps.""" 

243 

244 key_t = None 

245 changes = None 

246 defaults = None 

247 maps = None 

248 _observers = [] 

249 

250 def __init__(self, *maps, **kwargs): 

251 # type: (*Mapping, **Any) -> None 

252 maps = list(maps or [{}]) 

253 self.__dict__.update( 

254 key_t=kwargs.get('key_t'), 

255 maps=maps, 

256 changes=maps[0], 

257 defaults=maps[1:], 

258 ) 

259 

260 def add_defaults(self, d): 

261 # type: (Mapping) -> None 

262 d = force_mapping(d) 

263 self.defaults.insert(0, d) 

264 self.maps.insert(1, d) 

265 

266 def pop(self, key, *default): 

267 # type: (Any, *Any) -> Any 

268 try: 

269 return self.maps[0].pop(key, *default) 

270 except KeyError: 

271 raise KeyError( 

272 'Key not found in the first mapping: {!r}'.format(key)) 

273 

274 def __missing__(self, key): 

275 # type: (Any) -> Any 

276 raise KeyError(key) 

277 

278 def _key(self, key): 

279 # type: (Any) -> Any 

280 return self.key_t(key) if self.key_t is not None else key 

281 

282 def __getitem__(self, key): 

283 # type: (Any) -> Any 

284 _key = self._key(key) 

285 for mapping in self.maps: 

286 try: 

287 return mapping[_key] 

288 except KeyError: 

289 pass 

290 return self.__missing__(key) 

291 

292 def __setitem__(self, key, value): 

293 # type: (Any, Any) -> None 

294 self.changes[self._key(key)] = value 

295 

296 def __delitem__(self, key): 

297 # type: (Any) -> None 

298 try: 

299 del self.changes[self._key(key)] 

300 except KeyError: 

301 raise KeyError('Key not found in first mapping: {0!r}'.format(key)) 

302 

303 def clear(self): 

304 # type: () -> None 

305 self.changes.clear() 

306 

307 def get(self, key, default=None): 

308 # type: (Any, Any) -> Any 

309 try: 

310 return self[self._key(key)] 

311 except KeyError: 

312 return default 

313 

314 def __len__(self): 

315 # type: () -> int 

316 return len(set().union(*self.maps)) 

317 

318 def __iter__(self): 

319 return self._iterate_keys() 

320 

321 def __contains__(self, key): 

322 # type: (Any) -> bool 

323 key = self._key(key) 

324 return any(key in m for m in self.maps) 

325 

326 def __bool__(self): 

327 # type: () -> bool 

328 return any(self.maps) 

329 __nonzero__ = __bool__ # Py2 

330 

331 def setdefault(self, key, default=None): 

332 # type: (Any, Any) -> None 

333 key = self._key(key) 

334 if key not in self: 

335 self[key] = default 

336 

337 def update(self, *args, **kwargs): 

338 # type: (*Any, **Any) -> Any 

339 result = self.changes.update(*args, **kwargs) 

340 for callback in self._observers: 

341 callback(*args, **kwargs) 

342 return result 

343 

344 def __repr__(self): 

345 # type: () -> str 

346 return '{0.__class__.__name__}({1})'.format( 

347 self, ', '.join(map(repr, self.maps))) 

348 

349 @classmethod 

350 def fromkeys(cls, iterable, *args): 

351 # type: (type, Iterable, *Any) -> 'ChainMap' 

352 """Create a ChainMap with a single dict created from the iterable.""" 

353 return cls(dict.fromkeys(iterable, *args)) 

354 

355 def copy(self): 

356 # type: () -> 'ChainMap' 

357 return self.__class__(self.maps[0].copy(), *self.maps[1:]) 

358 __copy__ = copy # Py2 

359 

360 def _iter(self, op): 

361 # type: (Callable) -> Iterable 

362 # defaults must be first in the stream, so values in 

363 # changes take precedence. 

364 # pylint: disable=bad-reversed-sequence 

365 # Someone should teach pylint about properties. 

366 return chain(*[op(d) for d in reversed(self.maps)]) 

367 

368 def _iterate_keys(self): 

369 # type: () -> Iterable 

370 return uniq(self._iter(lambda d: d.keys())) 

371 iterkeys = _iterate_keys 

372 

373 def _iterate_items(self): 

374 # type: () -> Iterable 

375 return ((key, self[key]) for key in self) 

376 iteritems = _iterate_items 

377 

378 def _iterate_values(self): 

379 # type: () -> Iterable 

380 return (self[key] for key in self) 

381 itervalues = _iterate_values 

382 

383 def bind_to(self, callback): 

384 self._observers.append(callback) 

385 

386 if sys.version_info[0] == 3: # pragma: no cover 

387 keys = _iterate_keys 

388 items = _iterate_items 

389 values = _iterate_values 

390 

391 else: # noqa 

392 def keys(self): 

393 # type: () -> List[Any] 

394 return list(self._iterate_keys()) 

395 

396 def items(self): 

397 # type: () -> List[Tuple[Any, Any]] 

398 return list(self._iterate_items()) 

399 

400 def values(self): 

401 # type: () -> List[Any] 

402 return list(self._iterate_values()) 

403 

404 

405@python_2_unicode_compatible 

406class ConfigurationView(ChainMap, AttributeDictMixin): 

407 """A view over an applications configuration dictionaries. 

408 

409 Custom (but older) version of :class:`collections.ChainMap`. 

410 

411 If the key does not exist in ``changes``, the ``defaults`` 

412 dictionaries are consulted. 

413 

414 Arguments: 

415 changes (Mapping): Map of configuration changes. 

416 defaults (List[Mapping]): List of dictionaries containing 

417 the default configuration. 

418 """ 

419 

420 def __init__(self, changes, defaults=None, keys=None, prefix=None): 

421 # type: (Mapping, Mapping, List[str], str) -> None 

422 defaults = [] if defaults is None else defaults 

423 super(ConfigurationView, self).__init__(changes, *defaults) 

424 self.__dict__.update( 

425 prefix=prefix.rstrip('_') + '_' if prefix else prefix, 

426 _keys=keys, 

427 ) 

428 

429 def _to_keys(self, key): 

430 # type: (str) -> Sequence[str] 

431 prefix = self.prefix 

432 if prefix: 

433 pkey = prefix + key if not key.startswith(prefix) else key 

434 return match_case(pkey, prefix), key 

435 return key, 

436 

437 def __getitem__(self, key): 

438 # type: (str) -> Any 

439 keys = self._to_keys(key) 

440 getitem = super(ConfigurationView, self).__getitem__ 

441 for k in keys + ( 

442 tuple(f(key) for f in self._keys) if self._keys else ()): 

443 try: 

444 return getitem(k) 

445 except KeyError: 

446 pass 

447 try: 

448 # support subclasses implementing __missing__ 

449 return self.__missing__(key) 

450 except KeyError: 

451 if len(keys) > 1: 

452 raise KeyError( 

453 'Key not found: {0!r} (with prefix: {0!r})'.format(*keys)) 

454 raise 

455 

456 def __setitem__(self, key, value): 

457 # type: (str, Any) -> Any 

458 self.changes[self._key(key)] = value 

459 

460 def first(self, *keys): 

461 # type: (*str) -> Any 

462 return first(None, (self.get(key) for key in keys)) 

463 

464 def get(self, key, default=None): 

465 # type: (str, Any) -> Any 

466 try: 

467 return self[key] 

468 except KeyError: 

469 return default 

470 

471 def clear(self): 

472 # type: () -> None 

473 """Remove all changes, but keep defaults.""" 

474 self.changes.clear() 

475 

476 def __contains__(self, key): 

477 # type: (str) -> bool 

478 keys = self._to_keys(key) 

479 return any(any(k in m for k in keys) for m in self.maps) 

480 

481 def swap_with(self, other): 

482 # type: (ConfigurationView) -> None 

483 changes = other.__dict__['changes'] 

484 defaults = other.__dict__['defaults'] 

485 self.__dict__.update( 

486 changes=changes, 

487 defaults=defaults, 

488 key_t=other.__dict__['key_t'], 

489 prefix=other.__dict__['prefix'], 

490 maps=[changes] + defaults 

491 ) 

492 

493 

494@python_2_unicode_compatible 

495class LimitedSet(object): 

496 """Kind-of Set (or priority queue) with limitations. 

497 

498 Good for when you need to test for membership (`a in set`), 

499 but the set should not grow unbounded. 

500 

501 ``maxlen`` is enforced at all times, so if the limit is reached 

502 we'll also remove non-expired items. 

503 

504 You can also configure ``minlen``: this is the minimal residual size 

505 of the set. 

506 

507 All arguments are optional, and no limits are enabled by default. 

508 

509 Arguments: 

510 maxlen (int): Optional max number of items. 

511 Adding more items than ``maxlen`` will result in immediate 

512 removal of items sorted by oldest insertion time. 

513 

514 expires (float): TTL for all items. 

515 Expired items are purged as keys are inserted. 

516 

517 minlen (int): Minimal residual size of this set. 

518 .. versionadded:: 4.0 

519 

520 Value must be less than ``maxlen`` if both are configured. 

521 

522 Older expired items will be deleted, only after the set 

523 exceeds ``minlen`` number of items. 

524 

525 data (Sequence): Initial data to initialize set with. 

526 Can be an iterable of ``(key, value)`` pairs, 

527 a dict (``{key: insertion_time}``), or another instance 

528 of :class:`LimitedSet`. 

529 

530 Example: 

531 >>> s = LimitedSet(maxlen=50000, expires=3600, minlen=4000) 

532 >>> for i in range(60000): 

533 ... s.add(i) 

534 ... s.add(str(i)) 

535 ... 

536 >>> 57000 in s # last 50k inserted values are kept 

537 True 

538 >>> '10' in s # '10' did expire and was purged from set. 

539 False 

540 >>> len(s) # maxlen is reached 

541 50000 

542 >>> s.purge(now=monotonic() + 7200) # clock + 2 hours 

543 >>> len(s) # now only minlen items are cached 

544 4000 

545 >>>> 57000 in s # even this item is gone now 

546 False 

547 """ 

548 

549 max_heap_percent_overload = 15 

550 

551 def __init__(self, maxlen=0, expires=0, data=None, minlen=0): 

552 # type: (int, float, Mapping, int) -> None 

553 self.maxlen = 0 if maxlen is None else maxlen 

554 self.minlen = 0 if minlen is None else minlen 

555 self.expires = 0 if expires is None else expires 

556 self._data = {} 

557 self._heap = [] 

558 

559 if data: 

560 # import items from data 

561 self.update(data) 

562 

563 if not self.maxlen >= self.minlen >= 0: 

564 raise ValueError( 

565 'minlen must be a positive number, less or equal to maxlen.') 

566 if self.expires < 0: 

567 raise ValueError('expires cannot be negative!') 

568 

569 def _refresh_heap(self): 

570 # type: () -> None 

571 """Time consuming recreating of heap. Don't run this too often.""" 

572 self._heap[:] = [entry for entry in values(self._data)] 

573 heapify(self._heap) 

574 

575 def _maybe_refresh_heap(self): 

576 # type: () -> None 

577 if self._heap_overload >= self.max_heap_percent_overload: 

578 self._refresh_heap() 

579 

580 def clear(self): 

581 # type: () -> None 

582 """Clear all data, start from scratch again.""" 

583 self._data.clear() 

584 self._heap[:] = [] 

585 

586 def add(self, item, now=None): 

587 # type: (Any, float) -> None 

588 """Add a new item, or reset the expiry time of an existing item.""" 

589 now = now or monotonic() 

590 if item in self._data: 

591 self.discard(item) 

592 entry = (now, item) 

593 self._data[item] = entry 

594 heappush(self._heap, entry) 

595 if self.maxlen and len(self._data) >= self.maxlen: 

596 self.purge() 

597 

598 def update(self, other): 

599 # type: (Iterable) -> None 

600 """Update this set from other LimitedSet, dict or iterable.""" 

601 if not other: 

602 return 

603 if isinstance(other, LimitedSet): 

604 self._data.update(other._data) 

605 self._refresh_heap() 

606 self.purge() 

607 elif isinstance(other, dict): 

608 # revokes are sent as a dict 

609 for key, inserted in items(other): 

610 if isinstance(inserted, (tuple, list)): 

611 # in case someone uses ._data directly for sending update 

612 inserted = inserted[0] 

613 if not isinstance(inserted, float): 

614 raise ValueError( 

615 'Expecting float timestamp, got type ' 

616 '{0!r} with value: {1}'.format( 

617 type(inserted), inserted)) 

618 self.add(key, inserted) 

619 else: 

620 # XXX AVOID THIS, it could keep old data if more parties 

621 # exchange them all over and over again 

622 for obj in other: 

623 self.add(obj) 

624 

625 def discard(self, item): 

626 # type: (Any) -> None 

627 # mark an existing item as removed. If KeyError is not found, pass. 

628 self._data.pop(item, None) 

629 self._maybe_refresh_heap() 

630 pop_value = discard 

631 

632 def purge(self, now=None): 

633 # type: (float) -> None 

634 """Check oldest items and remove them if needed. 

635 

636 Arguments: 

637 now (float): Time of purging -- by default right now. 

638 This can be useful for unit testing. 

639 """ 

640 now = now or monotonic() 

641 now = now() if isinstance(now, Callable) else now 

642 if self.maxlen: 

643 while len(self._data) > self.maxlen: 

644 self.pop() 

645 # time based expiring: 

646 if self.expires: 

647 while len(self._data) > self.minlen >= 0: 

648 inserted_time, _ = self._heap[0] 

649 if inserted_time + self.expires > now: 

650 break # oldest item hasn't expired yet 

651 self.pop() 

652 

653 def pop(self, default=None): 

654 # type: (Any) -> Any 

655 """Remove and return the oldest item, or :const:`None` when empty.""" 

656 while self._heap: 

657 _, item = heappop(self._heap) 

658 try: 

659 self._data.pop(item) 

660 except KeyError: 

661 pass 

662 else: 

663 return item 

664 return default 

665 

666 def as_dict(self): 

667 # type: () -> Dict 

668 """Whole set as serializable dictionary. 

669 

670 Example: 

671 >>> s = LimitedSet(maxlen=200) 

672 >>> r = LimitedSet(maxlen=200) 

673 >>> for i in range(500): 

674 ... s.add(i) 

675 ... 

676 >>> r.update(s.as_dict()) 

677 >>> r == s 

678 True 

679 """ 

680 return {key: inserted for inserted, key in values(self._data)} 

681 

682 def __eq__(self, other): 

683 # type: (Any) -> bool 

684 return self._data == other._data 

685 

686 def __ne__(self, other): 

687 # type: (Any) -> bool 

688 return not self.__eq__(other) 

689 

690 def __repr__(self): 

691 # type: () -> str 

692 return REPR_LIMITED_SET.format( 

693 self, name=type(self).__name__, size=len(self), 

694 ) 

695 

696 def __iter__(self): 

697 # type: () -> Iterable 

698 return (i for _, i in sorted(values(self._data))) 

699 

700 def __len__(self): 

701 # type: () -> int 

702 return len(self._data) 

703 

704 def __contains__(self, key): 

705 # type: (Any) -> bool 

706 return key in self._data 

707 

708 def __reduce__(self): 

709 # type: () -> Any 

710 return self.__class__, ( 

711 self.maxlen, self.expires, self.as_dict(), self.minlen) 

712 

713 def __bool__(self): 

714 # type: () -> bool 

715 return bool(self._data) 

716 __nonzero__ = __bool__ # Py2 

717 

718 @property 

719 def _heap_overload(self): 

720 # type: () -> float 

721 """Compute how much is heap bigger than data [percents].""" 

722 return len(self._heap) * 100 / max(len(self._data), 1) - 100 

723 

724 

725MutableSet.register(LimitedSet) # noqa: E305 

726 

727 

728class Evictable(object): 

729 """Mixin for classes supporting the ``evict`` method.""" 

730 

731 Empty = Empty 

732 

733 def evict(self): 

734 # type: () -> None 

735 """Force evict until maxsize is enforced.""" 

736 self._evict(range=count) 

737 

738 def _evict(self, limit=100, range=range): 

739 # type: (int) -> None 

740 try: 

741 [self._evict1() for _ in range(limit)] 

742 except IndexError: 

743 pass 

744 

745 def _evict1(self): 

746 # type: () -> None 

747 if self._evictcount <= self.maxsize: 

748 raise IndexError() 

749 try: 

750 self._pop_to_evict() 

751 except self.Empty: 

752 raise IndexError() 

753 

754 

755@python_2_unicode_compatible 

756class Messagebuffer(Evictable): 

757 """A buffer of pending messages.""" 

758 

759 Empty = Empty 

760 

761 def __init__(self, maxsize, iterable=None, deque=deque): 

762 # type: (int, Iterable, Any) -> None 

763 self.maxsize = maxsize 

764 self.data = deque(iterable or []) 

765 self._append = self.data.append 

766 self._pop = self.data.popleft 

767 self._len = self.data.__len__ 

768 self._extend = self.data.extend 

769 

770 def put(self, item): 

771 # type: (Any) -> None 

772 self._append(item) 

773 self.maxsize and self._evict() 

774 

775 def extend(self, it): 

776 # type: (Iterable) -> None 

777 self._extend(it) 

778 self.maxsize and self._evict() 

779 

780 def take(self, *default): 

781 # type: (*Any) -> Any 

782 try: 

783 return self._pop() 

784 except IndexError: 

785 if default: 

786 return default[0] 

787 raise self.Empty() 

788 

789 def _pop_to_evict(self): 

790 # type: () -> None 

791 return self.take() 

792 

793 def __repr__(self): 

794 # type: () -> str 

795 return '<{0}: {1}/{2}>'.format( 

796 type(self).__name__, len(self), self.maxsize, 

797 ) 

798 

799 def __iter__(self): 

800 # type: () -> Iterable 

801 while 1: 

802 try: 

803 yield self._pop() 

804 except IndexError: 

805 break 

806 

807 def __len__(self): 

808 # type: () -> int 

809 return self._len() 

810 

811 def __contains__(self, item): 

812 # type: () -> bool 

813 return item in self.data 

814 

815 def __reversed__(self): 

816 # type: () -> Iterable 

817 return reversed(self.data) 

818 

819 def __getitem__(self, index): 

820 # type: (Any) -> Any 

821 return self.data[index] 

822 

823 @property 

824 def _evictcount(self): 

825 # type: () -> int 

826 return len(self) 

827 

828 

829Sequence.register(Messagebuffer) # noqa: E305 

830 

831 

832@python_2_unicode_compatible 

833class BufferMap(OrderedDict, Evictable): 

834 """Map of buffers.""" 

835 

836 Buffer = Messagebuffer 

837 Empty = Empty 

838 

839 maxsize = None 

840 total = 0 

841 bufmaxsize = None 

842 

843 def __init__(self, maxsize, iterable=None, bufmaxsize=1000): 

844 # type: (int, Iterable, int) -> None 

845 super(BufferMap, self).__init__() 

846 self.maxsize = maxsize 

847 self.bufmaxsize = 1000 

848 if iterable: 

849 self.update(iterable) 

850 self.total = sum(len(buf) for buf in items(self)) 

851 

852 def put(self, key, item): 

853 # type: (Any, Any) -> None 

854 self._get_or_create_buffer(key).put(item) 

855 self.total += 1 

856 self.move_to_end(key) # least recently used. 

857 self.maxsize and self._evict() 

858 

859 def extend(self, key, it): 

860 # type: (Any, Iterable) -> None 

861 self._get_or_create_buffer(key).extend(it) 

862 self.total += len(it) 

863 self.maxsize and self._evict() 

864 

865 def take(self, key, *default): 

866 # type: (Any, *Any) -> Any 

867 item, throw = None, False 

868 try: 

869 buf = self[key] 

870 except KeyError: 

871 throw = True 

872 else: 

873 try: 

874 item = buf.take() 

875 self.total -= 1 

876 except self.Empty: 

877 throw = True 

878 else: 

879 self.move_to_end(key) # mark as LRU 

880 

881 if throw: 

882 if default: 

883 return default[0] 

884 raise self.Empty() 

885 return item 

886 

887 def _get_or_create_buffer(self, key): 

888 # type: (Any) -> Messagebuffer 

889 try: 

890 return self[key] 

891 except KeyError: 

892 buf = self[key] = self._new_buffer() 

893 return buf 

894 

895 def _new_buffer(self): 

896 # type: () -> Messagebuffer 

897 return self.Buffer(maxsize=self.bufmaxsize) 

898 

899 def _LRUpop(self, *default): 

900 # type: (*Any) -> Any 

901 return self[self._LRUkey()].take(*default) 

902 

903 def _pop_to_evict(self): 

904 # type: () -> None 

905 for _ in range(100): 

906 key = self._LRUkey() 

907 buf = self[key] 

908 try: 

909 buf.take() 

910 except (IndexError, self.Empty): 

911 # buffer empty, remove it from mapping. 

912 self.pop(key) 

913 else: 

914 # we removed one item 

915 self.total -= 1 

916 # if buffer is empty now, remove it from mapping. 

917 if not len(buf): 

918 self.pop(key) 

919 else: 

920 # move to least recently used. 

921 self.move_to_end(key) 

922 break 

923 

924 def __repr__(self): 

925 # type: () -> str 

926 return '<{0}: {1}/{2}>'.format( 

927 type(self).__name__, self.total, self.maxsize, 

928 ) 

929 

930 @property 

931 def _evictcount(self): 

932 # type: () -> int 

933 return self.total