class HalfSnap: def __init__(self, bids: bool): if bids: self.data = SortedKeyList(key=lambda val: -val[0]) else: self.data = SortedKeyList(key=lambda val: val[0]) self.is_bids = bids self.time = None def fill(self, source): self.data.clear() for item in source: self.add(item) def add(self, item): price = item[0] size = item[1] self.data.add([price, size]) def update(self, price: float, size: float): key = -price if self.is_bids else price i = self.data.bisect_key_left(key) if 0 <= i < len(self.data): value = self.data[i] else: if size <= VERY_SMALL_NUMBER: return False self.data.add([price, size]) return True if size <= VERY_SMALL_NUMBER: if value[0] == price: self.data.discard(value) return True else: return False if value[0] == price: self.data[i][1] = size else: self.data.add([price, size]) return True def delete(self, price: float): return self.updatef(price, 0.0)
def test_remove(): slt = SortedKeyList(key=modulo) assert slt.discard(0) == None assert len(slt) == 0 slt._check() slt = SortedKeyList([1, 2, 2, 2, 3, 3, 5], key=modulo) slt._reset(4) slt.remove(2) slt._check() assert all(tup[0] == tup[1] for tup in zip(slt, [1, 2, 2, 3, 3, 5]))
def test_discard(): slt = SortedKeyList(key=negate) assert slt.discard(0) == None assert len(slt) == 0 slt._check() slt = SortedKeyList([1, 2, 2, 2, 3, 3, 5], key=negate) slt._reset(4) slt.discard(6) slt._check() slt.discard(4) slt._check() slt.discard(2) slt._check() assert all(tup[0] == tup[1] for tup in zip(slt, reversed([1, 2, 2, 3, 3, 5])))
class SortedIntvls: """ """ def __init__(self): # we sort by increasing start offset then increasing annotation id for this self._by_start = SortedKeyList(key=lambda x: (x[0], x[2])) # for this we sort by end offset only self._by_end = SortedKeyList(key=lambda x: x[1]) def add(self, start, end, data): """ Adds an interval. """ self._by_start.add((start, end, data)) self._by_end.add((start, end, data)) def update(self, tupleiterable): """ Updates from an iterable of intervals. """ self._by_start.update(tupleiterable) self._by_end.update(tupleiterable) def remove(self, start, end, data): """ Removes an interval, exception if the interval does not exist. """ self._by_start.remove((start, end, data)) self._by_end.remove((start, end, data)) def discard(self, start, end, data): """ Removes and interval, do nothing if the interval does not exist. """ self._by_start.discard((start, end, data)) self._by_end.discard((start, end, data)) def __len__(self): """ Returns the number of intervals. """ return len(self._by_start) def starting_at(self, offset): """ Returns an iterable of (start, end, data) tuples where start==offset """ return self._by_start.irange_key(min_key=(offset, 0), max_key=(offset, sys.maxsize)) def ending_at(self, offset): """ Returns an iterable of (start, end, data) tuples where end==offset """ return self._by_end.irange_key(min_key=offset, max_key=offset) def at(self, start, end): """ Returns an iterable of tuples where start==start and end==end """ for intvl in self._by_start.irange_key(min_key=(start, 0), max_key=(start, sys.maxsize)): if intvl[1] == end: yield intvl def within(self, start, end): """ Returns intervals which are fully contained within start...end """ # get all the intervals that start within the range, then keep those which also end within the range for intvl in self._by_start.irange_key(min_key=(start, 0), max_key=(end, sys.maxsize)): if intvl[1] <= end: yield intvl def starting_from(self, offset): """ Returns intervals that start at or after offset. """ return self._by_start.irange_key(min_key=(offset, 0)) def starting_before(self, offset): """ Returns intervals that start before offset. """ return self._by_start.irange_key(max_key=(offset - 1, sys.maxsize)) def ending_to(self, offset): """ Returns intervals that end before or at the given end offset. """ return self._by_end.irange_key(max_key=offset) def ending_after(self, offset): """ Returns intervals the end after the given offset. """ return self._by_end.irange_key(min_key=offset + 1) def covering(self, start, end): """ Returns intervals that contain the given range. """ # All intervals that start at or before the start and end at or after the end offset # we do this by first getting the intervals the start before or atthe start # then filtering by end for intvl in self._by_start.irange_key(max_key=(start, sys.maxsize)): if intvl[1] >= end: yield intvl def overlapping(self, start, end): """ Returns intervals that overlap with the given range. """ # All intervals where the start or end offset lies within the given range. # This excludes the ones where the end offset is before the start or # where the start offset is after the end of the range. # Here we do this by looking at all intervals where the start offset is before the # end of the range. This still includes those which also end before the start of the range # so we check in addition that the end is larger than the start of the range. for intvl in self._by_start.irange_key(max_key=(end - 1, sys.maxsize)): if intvl[1] > start + 1: yield intvl def firsts(self): """ Yields all intervals which start at the smallest known offset. """ laststart = None # logger.info("DEBUG: set laststart to None") for intvl in self._by_start.irange_key(): # logger.info("DEBUG: checking interval {}".format(intvl)) if laststart is None: laststart = intvl[0] # logger.info("DEBUG: setting laststart to {} and yielding {}".format(intvl[0], intvl)) yield intvl elif intvl[0] == laststart: # logger.info("DEBUG: yielding {}".format(intvl)) yield intvl else: # logger.info("DEBUG: returning since we got {}".format(intvl)) return def lasts(self): """ Yields all intervals which start at the last known start offset. """ laststart = None for intvl in reversed(self._by_start): if laststart is None: laststart = intvl[0] yield intvl elif intvl[0] == laststart: yield intvl else: return def min_start(self): """ Returns the smallest known start offset. """ return self._by_start[0][0] def max_end(self): """ Returns the biggest known end offset. """ return self._by_end[-1][1] def irange(self, minoff=None, maxoff=None, reverse=False, inclusive=(True, True)): """ Yields an iterator of intervals with a start offset between minoff and maxoff, inclusive. Args: minoff: minimum offset, default None indicates any maxoff: maximum offset, default None indicates any reverse: if `True` yield in reverse order inclusive: if the minoff and maxoff values should be inclusive, default is (True,True) Returns: """ return self._by_start.irange_key(min_key=minoff, max_key=maxoff, reverse=reverse, inclusive=inclusive) def __repr__(self): return "SortedIntvls({},{})".format(self._by_start, self._by_end)