def bracket(iterable, origin, interval_size, already_sorted=False, intervals_right_closed=False, coalesce=False): if not already_sorted: sorted_indices, iterable = zip( *sorted([(i, v) for i, v in enumerate(iterable)], key=operator.itemgetter(1))) brackets = [] bracket_indices = [] interval_offset = None interval_left = None interval_right = None for x in iterable: if interval_offset is None or x - interval_left >= interval_size: new_interval_offset = (x - origin) // interval_size new_interval_left = origin + new_interval_offset * interval_size if intervals_right_closed and x == new_interval_left: new_interval_offset -= 1 new_interval_left -= interval_size if coalesce and (interval_offset is not None) and ( new_interval_left <= brackets[-1].interval.right): interval_right = new_interval_left + interval_size brackets[-1].interval = brackets[-1].interval.replace_right( interval_right) elif interval_offset is None or new_interval_offset != interval_offset: interval_offset = new_interval_offset interval_left = new_interval_left interval_right = interval_left + interval_size brackets.append( Bracket( intervals.Interval(interval_left, interval_right, not intervals_right_closed, intervals_right_closed), interval_offset)) bracket_indices.append(len(brackets) - 1) if not already_sorted: new_bracket_indices = [None] * len(bracket_indices) for i in range(len(bracket_indices)): new_bracket_indices[sorted_indices[i]] = bracket_indices[i] bracket_indices = new_bracket_indices return brackets, bracket_indices
def test_intervals(self): result = utils.intervals(start=0, end=15, delta=5, intervals_right_closed=False) self.assertSequenceEqual(result, [ intervals.Interval(0, 5, left_closed=True, right_closed=False), intervals.Interval(5, 10, left_closed=True, right_closed=False), intervals.Interval(10, 15, left_closed=True, right_closed=False) ]) result = utils.intervals(start=0, end=15, delta=5, intervals_right_closed=True) self.assertSequenceEqual(result, [ intervals.Interval(0, 5, left_closed=False, right_closed=True), intervals.Interval(5, 10, left_closed=False, right_closed=True), intervals.Interval(10, 15, left_closed=False, right_closed=True) ]) result = utils.intervals(start=0, end=15, delta=4, intervals_right_closed=False) self.assertSequenceEqual(result, [ intervals.Interval(0, 4, left_closed=True, right_closed=False), intervals.Interval(4, 8, left_closed=True, right_closed=False), intervals.Interval(8, 12, left_closed=True, right_closed=False), intervals.Interval(12, 15, left_closed=True, right_closed=False) ]) result = utils.intervals(start=0, end=15, delta=4, intervals_right_closed=True) self.assertSequenceEqual(result, [ intervals.Interval(0, 4, left_closed=False, right_closed=True), intervals.Interval(4, 8, left_closed=False, right_closed=True), intervals.Interval(8, 12, left_closed=False, right_closed=True), intervals.Interval(12, 15, left_closed=False, right_closed=True) ]) result = utils.intervals(start=dt.date(2019, 8, 31), end=dt.date(2019, 9, 15), delta=dt.timedelta(days=5), intervals_right_closed=False) self.assertSequenceEqual(result, [ intervals.Interval(dt.date(2019, 8, 31), dt.date(2019, 9, 5), left_closed=True, right_closed=False), intervals.Interval(dt.date(2019, 9, 5), dt.date(2019, 9, 10), left_closed=True, right_closed=False), intervals.Interval(dt.date(2019, 9, 10), dt.date(2019, 9, 15), left_closed=True, right_closed=False) ]) result = utils.intervals(start=dt.date(2019, 8, 31), end=dt.date(2019, 9, 15), delta=dt.timedelta(days=5), intervals_right_closed=True) self.assertSequenceEqual(result, [ intervals.Interval(dt.date(2019, 8, 31), dt.date(2019, 9, 5), left_closed=False, right_closed=True), intervals.Interval(dt.date(2019, 9, 5), dt.date(2019, 9, 10), left_closed=False, right_closed=True), intervals.Interval(dt.date(2019, 9, 10), dt.date(2019, 9, 15), left_closed=False, right_closed=True) ]) result = utils.intervals(start=dt.date(2019, 8, 31), end=dt.date(2019, 9, 15), delta=dt.timedelta(days=4), intervals_right_closed=False) self.assertSequenceEqual(result, [ intervals.Interval(dt.date(2019, 8, 31), dt.date(2019, 9, 4), left_closed=True, right_closed=False), intervals.Interval(dt.date(2019, 9, 4), dt.date(2019, 9, 8), left_closed=True, right_closed=False), intervals.Interval(dt.date(2019, 9, 8), dt.date(2019, 9, 12), left_closed=True, right_closed=False), intervals.Interval(dt.date(2019, 9, 12), dt.date(2019, 9, 15), left_closed=True, right_closed=False) ]) result = utils.intervals(start=dt.date(2019, 8, 31), end=dt.date(2019, 9, 15), delta=dt.timedelta(days=4), intervals_right_closed=True) self.assertSequenceEqual(result, [ intervals.Interval(dt.date(2019, 8, 31), dt.date(2019, 9, 4), left_closed=False, right_closed=True), intervals.Interval(dt.date(2019, 9, 4), dt.date(2019, 9, 8), left_closed=False, right_closed=True), intervals.Interval(dt.date(2019, 9, 8), dt.date(2019, 9, 12), left_closed=False, right_closed=True), intervals.Interval(dt.date(2019, 9, 12), dt.date(2019, 9, 15), left_closed=False, right_closed=True) ]) result = utils.intervals(start=dt.datetime(2019, 10, 8, 0), end=dt.datetime(2019, 10, 8, 15), delta=dt.timedelta(hours=5), intervals_right_closed=False) self.assertSequenceEqual(result, [ intervals.Interval(dt.datetime(2019, 10, 8, 0), dt.datetime(2019, 10, 8, 5), left_closed=True, right_closed=False), intervals.Interval(dt.datetime(2019, 10, 8, 5), dt.datetime(2019, 10, 8, 10), left_closed=True, right_closed=False), intervals.Interval(dt.datetime(2019, 10, 8, 10), dt.datetime(2019, 10, 8, 15), left_closed=True, right_closed=False) ]) result = utils.intervals(start=dt.datetime(2019, 10, 8, 0), end=dt.datetime(2019, 10, 8, 15), delta=dt.timedelta(hours=5), intervals_right_closed=True) self.assertSequenceEqual(result, [ intervals.Interval(dt.datetime(2019, 10, 8, 0), dt.datetime(2019, 10, 8, 5), left_closed=False, right_closed=True), intervals.Interval(dt.datetime(2019, 10, 8, 5), dt.datetime(2019, 10, 8, 10), left_closed=False, right_closed=True), intervals.Interval(dt.datetime(2019, 10, 8, 10), dt.datetime(2019, 10, 8, 15), left_closed=False, right_closed=True) ]) result = utils.intervals(start=dt.datetime(2019, 10, 8, 0), end=dt.datetime(2019, 10, 8, 15), delta=dt.timedelta(hours=4), intervals_right_closed=False) self.assertSequenceEqual(result, [ intervals.Interval(dt.datetime(2019, 10, 8, 0), dt.datetime(2019, 10, 8, 4), left_closed=True, right_closed=False), intervals.Interval(dt.datetime(2019, 10, 8, 4), dt.datetime(2019, 10, 8, 8), left_closed=True, right_closed=False), intervals.Interval(dt.datetime(2019, 10, 8, 8), dt.datetime(2019, 10, 8, 12), left_closed=True, right_closed=False), intervals.Interval(dt.datetime(2019, 10, 8, 12), dt.datetime(2019, 10, 8, 15), left_closed=True, right_closed=False) ]) result = utils.intervals(start=dt.datetime(2019, 10, 8, 0), end=dt.datetime(2019, 10, 8, 15), delta=dt.timedelta(hours=4), intervals_right_closed=True) self.assertSequenceEqual(result, [ intervals.Interval(dt.datetime(2019, 10, 8, 0), dt.datetime(2019, 10, 8, 4), left_closed=False, right_closed=True), intervals.Interval(dt.datetime(2019, 10, 8, 4), dt.datetime(2019, 10, 8, 8), left_closed=False, right_closed=True), intervals.Interval(dt.datetime(2019, 10, 8, 8), dt.datetime(2019, 10, 8, 12), left_closed=False, right_closed=True), intervals.Interval(dt.datetime(2019, 10, 8, 12), dt.datetime(2019, 10, 8, 15), left_closed=False, right_closed=True) ])
def test_bracket(self): data = [8, 11, 12, 13, 14, 27, 29, 37, 49, 50, 51, 79, 85] brackets, bracket_indices = utils.bracket(data, 3, 5) self.assertSequenceEqual(brackets, [ utils.Bracket(intervals.Interval(8, 13, left_closed=True), 1), utils.Bracket(intervals.Interval(13, 18, left_closed=True), 2), utils.Bracket(intervals.Interval(23, 28, left_closed=True), 4), utils.Bracket(intervals.Interval(28, 33, left_closed=True), 5), utils.Bracket(intervals.Interval(33, 38, left_closed=True), 6), utils.Bracket(intervals.Interval(48, 53, left_closed=True), 9), utils.Bracket(intervals.Interval(78, 83, left_closed=True), 15), utils.Bracket(intervals.Interval(83, 88, left_closed=True), 16) ]) self.assertSequenceEqual(bracket_indices, [0, 0, 0, 1, 1, 2, 3, 4, 5, 5, 5, 6, 7]) brackets, bracket_indices = utils.bracket(data, 3, 5, intervals_right_closed=True) self.assertSequenceEqual(brackets, [ utils.Bracket(intervals.Interval(3, 8, right_closed=True), 0), utils.Bracket(intervals.Interval(8, 13, right_closed=True), 1), utils.Bracket(intervals.Interval(13, 18, right_closed=True), 2), utils.Bracket(intervals.Interval(23, 28, right_closed=True), 4), utils.Bracket(intervals.Interval(28, 33, right_closed=True), 5), utils.Bracket(intervals.Interval(33, 38, right_closed=True), 6), utils.Bracket(intervals.Interval(48, 53, right_closed=True), 9), utils.Bracket(intervals.Interval(78, 83, right_closed=True), 15), utils.Bracket(intervals.Interval(83, 88, right_closed=True), 16) ]) self.assertSequenceEqual(bracket_indices, [0, 1, 1, 1, 2, 3, 4, 5, 6, 6, 6, 7, 8]) brackets, bracket_indices = utils.bracket(data, 3, 5, coalesce=True) self.assertSequenceEqual(brackets, [ utils.Bracket(intervals.Interval(8, 18, left_closed=True), 1), utils.Bracket(intervals.Interval(23, 38, left_closed=True), 4), utils.Bracket(intervals.Interval(48, 53, left_closed=True), 9), utils.Bracket(intervals.Interval(78, 88, left_closed=True), 15) ]) self.assertSequenceEqual(bracket_indices, [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3]) brackets, bracket_indices = utils.bracket(data, 3, 5, intervals_right_closed=True, coalesce=True) self.assertSequenceEqual(brackets, [ utils.Bracket(intervals.Interval(3, 18, right_closed=True), 0), utils.Bracket(intervals.Interval(23, 38, right_closed=True), 4), utils.Bracket(intervals.Interval(48, 53, right_closed=True), 9), utils.Bracket(intervals.Interval(78, 88, right_closed=True), 15) ]) self.assertSequenceEqual(bracket_indices, [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3]) data = [ dt.date(2017, 1, 31) + dt.timedelta(days=x) for x in [8, 11, 12, 13, 14, 27, 29, 37, 49, 50, 51, 79, 85] ] brackets, bracket_indices = utils.bracket(data, dt.date(2017, 2, 3), dt.timedelta(days=5)) self.assertSequenceEqual(brackets, [ utils.Bracket( intervals.Interval(dt.date(2017, 2, 8), dt.date(2017, 2, 13), left_closed=True), 1), utils.Bracket( intervals.Interval(dt.date(2017, 2, 13), dt.date(2017, 2, 18), left_closed=True), 2), utils.Bracket( intervals.Interval(dt.date(2017, 2, 23), dt.date(2017, 2, 28), left_closed=True), 4), utils.Bracket( intervals.Interval(dt.date(2017, 2, 28), dt.date(2017, 3, 5), left_closed=True), 5), utils.Bracket( intervals.Interval(dt.date(2017, 3, 5), dt.date(2017, 3, 10), left_closed=True), 6), utils.Bracket( intervals.Interval(dt.date(2017, 3, 20), dt.date(2017, 3, 25), left_closed=True), 9), utils.Bracket( intervals.Interval(dt.date(2017, 4, 19), dt.date(2017, 4, 24), left_closed=True), 15), utils.Bracket( intervals.Interval(dt.date(2017, 4, 24), dt.date(2017, 4, 29), left_closed=True), 16) ]) self.assertSequenceEqual(bracket_indices, [0, 0, 0, 1, 1, 2, 3, 4, 5, 5, 5, 6, 7]) brackets, bracket_indices = utils.bracket(data, dt.date(2017, 2, 3), dt.timedelta(days=5), intervals_right_closed=True) self.assertSequenceEqual(brackets, [ utils.Bracket( intervals.Interval(dt.date(2017, 2, 3), dt.date(2017, 2, 8), right_closed=True), 0), utils.Bracket( intervals.Interval(dt.date(2017, 2, 8), dt.date(2017, 2, 13), right_closed=True), 1), utils.Bracket( intervals.Interval(dt.date(2017, 2, 13), dt.date(2017, 2, 18), right_closed=True), 2), utils.Bracket( intervals.Interval(dt.date(2017, 2, 23), dt.date(2017, 2, 28), right_closed=True), 4), utils.Bracket( intervals.Interval(dt.date(2017, 2, 28), dt.date(2017, 3, 5), right_closed=True), 5), utils.Bracket( intervals.Interval(dt.date(2017, 3, 5), dt.date(2017, 3, 10), right_closed=True), 6), utils.Bracket( intervals.Interval(dt.date(2017, 3, 20), dt.date(2017, 3, 25), right_closed=True), 9), utils.Bracket( intervals.Interval(dt.date(2017, 4, 19), dt.date(2017, 4, 24), right_closed=True), 15), utils.Bracket( intervals.Interval(dt.date(2017, 4, 24), dt.date(2017, 4, 29), right_closed=True), 16) ]) self.assertSequenceEqual(bracket_indices, [0, 1, 1, 1, 2, 3, 4, 5, 6, 6, 6, 7, 8]) brackets, bracket_indices = utils.bracket(data, dt.date(2017, 2, 3), dt.timedelta(days=5), coalesce=True) self.assertSequenceEqual(brackets, [ utils.Bracket( intervals.Interval(dt.date(2017, 2, 8), dt.date(2017, 2, 18), left_closed=True), 1), utils.Bracket( intervals.Interval(dt.date(2017, 2, 23), dt.date(2017, 3, 10), left_closed=True), 4), utils.Bracket( intervals.Interval(dt.date(2017, 3, 20), dt.date(2017, 3, 25), left_closed=True), 9), utils.Bracket( intervals.Interval(dt.date(2017, 4, 19), dt.date(2017, 4, 29), left_closed=True), 15) ]) self.assertSequenceEqual(bracket_indices, [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3]) brackets, bracket_indices = utils.bracket(data, dt.date(2017, 2, 3), dt.timedelta(days=5), intervals_right_closed=True, coalesce=True) self.assertSequenceEqual(brackets, [ utils.Bracket( intervals.Interval(dt.date(2017, 2, 3), dt.date(2017, 2, 18), right_closed=True), 0), utils.Bracket( intervals.Interval(dt.date(2017, 2, 23), dt.date(2017, 3, 10), right_closed=True), 4), utils.Bracket( intervals.Interval(dt.date(2017, 3, 20), dt.date(2017, 3, 25), right_closed=True), 9), utils.Bracket( intervals.Interval(dt.date(2017, 4, 19), dt.date(2017, 4, 29), right_closed=True), 15) ]) self.assertSequenceEqual(bracket_indices, [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3]) data = [ dt.datetime(2017, 1, 31, 0, 0, 0) + dt.timedelta(minutes=x) for x in [8, 11, 12, 13, 14, 27, 29, 37, 49, 50, 51, 79, 85] ] brackets, bracket_indices = utils.bracket( data, dt.datetime(2017, 1, 31, 0, 3, 0), dt.timedelta(minutes=5)) self.assertSequenceEqual(brackets, [ utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 8), dt.datetime(2017, 1, 31, 0, 13), left_closed=True), 1), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 13), dt.datetime(2017, 1, 31, 0, 18), left_closed=True), 2), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 23), dt.datetime(2017, 1, 31, 0, 28), left_closed=True), 4), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 28), dt.datetime(2017, 1, 31, 0, 33), left_closed=True), 5), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 33), dt.datetime(2017, 1, 31, 0, 38), left_closed=True), 6), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 48), dt.datetime(2017, 1, 31, 0, 53), left_closed=True), 9), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 1, 18), dt.datetime(2017, 1, 31, 1, 23), left_closed=True), 15), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 1, 23), dt.datetime(2017, 1, 31, 1, 28), left_closed=True), 16) ]) self.assertSequenceEqual(bracket_indices, [0, 0, 0, 1, 1, 2, 3, 4, 5, 5, 5, 6, 7]) brackets, bracket_indices = utils.bracket(data, dt.datetime( 2017, 1, 31, 0, 3, 0), dt.timedelta(minutes=5), intervals_right_closed=True) self.assertSequenceEqual(brackets, [ utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 3), dt.datetime(2017, 1, 31, 0, 8), right_closed=True), 0), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 8), dt.datetime(2017, 1, 31, 0, 13), right_closed=True), 1), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 13), dt.datetime(2017, 1, 31, 0, 18), right_closed=True), 2), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 23), dt.datetime(2017, 1, 31, 0, 28), right_closed=True), 4), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 28), dt.datetime(2017, 1, 31, 0, 33), right_closed=True), 5), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 33), dt.datetime(2017, 1, 31, 0, 38), right_closed=True), 6), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 48), dt.datetime(2017, 1, 31, 0, 53), right_closed=True), 9), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 1, 18), dt.datetime(2017, 1, 31, 1, 23), right_closed=True), 15), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 1, 23), dt.datetime(2017, 1, 31, 1, 28), right_closed=True), 16) ]) self.assertSequenceEqual(bracket_indices, [0, 1, 1, 1, 2, 3, 4, 5, 6, 6, 6, 7, 8]) brackets, bracket_indices = utils.bracket(data, dt.datetime( 2017, 1, 31, 0, 3, 0), dt.timedelta(minutes=5), coalesce=True) self.assertSequenceEqual(brackets, [ utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 8), dt.datetime(2017, 1, 31, 0, 18), left_closed=True), 1), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 23), dt.datetime(2017, 1, 31, 0, 38), left_closed=True), 4), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 48), dt.datetime(2017, 1, 31, 0, 53), left_closed=True), 9), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 1, 18), dt.datetime(2017, 1, 31, 1, 28), left_closed=True), 15) ]) self.assertSequenceEqual(bracket_indices, [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3]) brackets, bracket_indices = utils.bracket(data, dt.datetime( 2017, 1, 31, 0, 3, 0), dt.timedelta(minutes=5), intervals_right_closed=True, coalesce=True) self.assertSequenceEqual(brackets, [ utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 3), dt.datetime(2017, 1, 31, 0, 18), right_closed=True), 0), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 23), dt.datetime(2017, 1, 31, 0, 38), right_closed=True), 4), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 0, 48), dt.datetime(2017, 1, 31, 0, 53), right_closed=True), 9), utils.Bracket( intervals.Interval(dt.datetime(2017, 1, 31, 1, 18), dt.datetime(2017, 1, 31, 1, 28), right_closed=True), 15) ]) self.assertSequenceEqual(bracket_indices, [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3])