diff --git a/CHANGELOG.md b/CHANGELOG.md index 670449e47b31cf03d9e68abc9d328ea9da2c2f51..3e67db7808ba475564e18702878fdd5dd22b8589 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,8 @@ +## v2.0.11 + +* dodano Heap + ## v2.0.10 * bugfix release diff --git a/satella/coding/__init__.py b/satella/coding/__init__.py index fa087969089763eca1fb0279e94710a35d571c08..bf971bff957f655b4955f17a016cdf7a1db44ac9 100644 --- a/satella/coding/__init__.py +++ b/satella/coding/__init__.py @@ -7,6 +7,6 @@ from __future__ import print_function, absolute_import, division from .typecheck import typed, List, Tuple, Dict, NewType, Callable, Sequence, \ TypeVar, Generic, Mapping, Iterable, Union, Any, Optional, CallSignature -from .structures import TimeBasedHeap, CallableGroup +from .structures import TimeBasedHeap, CallableGroup, Heap from .monitor import Monitor, RMonitor from .algos import merge_dicts diff --git a/satella/coding/structures.py b/satella/coding/structures.py index 9d9b825b76db947d43db2c7a5cebdca11924509b..2f84f4174e163a467d3bdd5f6a2861f1d1af4b6a 100644 --- a/satella/coding/structures.py +++ b/satella/coding/structures.py @@ -2,13 +2,19 @@ from __future__ import print_function, absolute_import, division import six import logging +import copy import heapq -from .typecheck import typed, Callable +import functools +from .typecheck import typed, Callable, Iterable logger = logging.getLogger(__name__) - +__all__ = [ + 'CallableGroup', + 'Heap', + 'TimeBasedHeap' +] class CallableGroup(object): @@ -81,7 +87,134 @@ class CallableGroup(object): return results -class TimeBasedHeap(object): +def _extras_to_one(fun): + @functools.wraps(fun) + def inner(self, a, *args): + return fun(self, ((a, ) + args) if len(args) > 0 else a) + return inner + + +class Heap(object): + """ + Sane heap as object - not like heapq. + + Goes from lowest-to-highest (first popped is smallest). + Standard Python comparision rules apply. + + Not thread-safe + """ + + __slots__ = ('heap', ) # this is rather private, plz + + # TODO needs tests + @typed(object, (None, Iterable)) + def __init__(self, from_list=None): + if from_list is None: + self.heap = [] + else: + self.heap = heapq.heapify(list(from_list)) + + # TODO needs tests + @_extras_to_one + def push(self, item): + """ + Use it like: + + heap.push(3) + + or: + + heap.push(4, myobject) + """ + heapq.heappush(self.heap, item) + + # TODO needs tests + def __deepcopy__(self): + h = Heap() + h.heap = copy.deepcopy(self.heap) + return h + + # TODO needs tests + def __copy__(self): + h = Heap() + h.heap = copy.copy(self.heap) + return h + + # TODO needs tests + def pop(self): + """ + :raises IndexError: on empty heap + """ + return heapq.heappop(self.heap) + + # TODO needs tests + @typed(object, Callable, Callable) + def filtermap(self, filterer=lambda i: True, mapfun=lambda i: i): + """ + Get only items that return True when condition(item) is True. Apply a transform: item' = item(condition) on + the rest. Maintain heap invariant. + """ + self.heap = [mapfun(s) for s in self.heap if filterer(s)] + heapq.heapify(self.heap) + + @typed(returns=bool) + def __bool__(self): + """ + Is this empty? + """ + return len(self.heap) > 0 + + @typed(returns=Iterable) + def iter_ascending(self): + """ + Return an iterator returning all elements in this heap sorted ascending. + State of the heap is not changed + :return: Iterator + """ + cph = self.copy() + while cph: + yield cph.pop() + + @typed(object, object, returns=Iterable) + def get_less_than(self, less): + """ + Return all elements less (sharp inequality) than particular value. + + This changes state of the heap + :param less: value to compare against + :return: Iterator + """ + while self: + if self.heap[0] < less: + return + yield self.pop() + + @typed(returns=Iterable) + def iter_descending(self): + """ + Return an iterator returning all elements in this heap sorted descending. + State of the heap is not changed + :return: Iterator + """ + return reversed(self.iter_ascending()) + + @typed(returns=six.integer_types) + def __len__(self): + return len(self) + + def __str__(self): + return '<satella.coding.Heap: %s elements>' % (len(self.heap, )) + + def __unicode__(self): + return six.text_type(str(self)) + + def __repr__(self): + return u'<satella.coding.Heap>' + + def __in__(self, item): + return item in self.heap + +class TimeBasedHeap(Heap): """ A heap of items sorted by timestamps. @@ -98,7 +231,7 @@ class TimeBasedHeap(object): """ Initialize an empty heap """ - self.heap = [] + super(TimeBasedHeap, self).__init__() @typed(None, (float, int), None) def put(self, timestamp, item): @@ -107,7 +240,7 @@ class TimeBasedHeap(object): :param timestamp: timestamp for this item :param item: object """ - heapq.heappush(self.heap, (timestamp, item)) + self.push(timestamp, item) @typed(None, (float, int)) def pop_less_than(self, timestamp): @@ -117,17 +250,11 @@ class TimeBasedHeap(object): Items will be removed from heap :return: list of tuple(timestamp::float, item) """ - out = [] - while len(self.heap) > 0: - if self.heap[0][0] >= timestamp: - return out - out.append(heapq.heappop(self.heap)) - return out + return list(self.get_less_than(timestamp)) def remove(self, item): """ Remove all things equal to item """ - self.heap = [q for q in self.heap if q != item] - heapq.heapify(self.heap) + self.filter(lambda i: i != item) diff --git a/setup.py b/setup.py index 804cddc3bf1d293d369256b3dd4a754632fceedb..478d348200388c724adae79f432e14e64559fac4 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ from setuptools import setup, find_packages setup(name='satella', - version='2.0.10', + version='2.0.11rc1', description=u'Utilities for writing servers in Python', author=u'Piotr MaĹlanka', author_email='piotrm@smok.co', diff --git a/tests/test_coding/test_structures.py b/tests/test_coding/test_structures.py index a16bff84ae73189d82f5199ed8e8cb5c65fba7ba..43baebb4c3f7f109fefd5e93b891979c5fa2cb3f 100644 --- a/tests/test_coding/test_structures.py +++ b/tests/test_coding/test_structures.py @@ -2,7 +2,7 @@ from __future__ import print_function, absolute_import, division import six import unittest -from satella.coding import TimeBasedHeap +from satella.coding import TimeBasedHeap, Heap class TestTimeBasedHeap(unittest.TestCase): @@ -19,3 +19,21 @@ class TestTimeBasedHeap(unittest.TestCase): self.assertIn((10, 'ala'), q) self.assertIn((20, 'ma'), q) self.assertNotIn((30, 'kota'), q) + + +class TestHeap(unittest.TestCase): + def test_tbh(self): + + tbh = Heap() + + tbh.put((10, 'ala')) + tbh.put(20, 'ma') + + self.assertIn((10, 'ala'), tbh) + self.assertIn((20, 'ma'), tbh) + + tbh.filtermap(lambda x: x[0] != 20, lambda x: x[0]+10, 'azomg') + + self.assertIn((20, 'azomg'), tbh) + self.assertNotIn((10, 'ala'), tbh) + self.assertNotIn((20, 'ma'), tbh)