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_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])))
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)
