mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-29 04:08:41 +08:00
Joe Jevnik
365 lines
8.4 KiB
Python
365 lines
8.4 KiB
Python
import operator as op
|
|
|
|
from six import PY2
|
|
from toolz import peek
|
|
|
|
from zipline.utils.functional import foldr
|
|
|
|
|
|
if PY2:
|
|
class range(object):
|
|
"""Lazy range object with constant time containment check.
|
|
|
|
The arguments are the same as ``range``.
|
|
"""
|
|
__slots__ = 'start', 'stop', 'step'
|
|
|
|
def __init__(self, stop, *args):
|
|
if len(args) > 2:
|
|
raise TypeError(
|
|
'range takes at most 3 arguments (%d given)' % len(args)
|
|
)
|
|
|
|
if not args:
|
|
self.start = 0
|
|
self.stop = stop
|
|
self.step = 1
|
|
else:
|
|
self.start = stop
|
|
self.stop = args[0]
|
|
try:
|
|
self.step = args[1]
|
|
except IndexError:
|
|
self.step = 1
|
|
|
|
if self.step == 0:
|
|
raise ValueError('range step must not be zero')
|
|
|
|
def __iter__(self):
|
|
"""
|
|
Examples
|
|
--------
|
|
>>> list(range(1))
|
|
[0]
|
|
>>> list(range(5))
|
|
[0, 1, 2, 3, 4]
|
|
>>> list(range(1, 5))
|
|
[1, 2, 3, 4]
|
|
>>> list(range(0, 5, 2))
|
|
[0, 2, 4]
|
|
>>> list(range(5, 0, -1))
|
|
[5, 4, 3, 2, 1]
|
|
>>> list(range(5, 0, 1))
|
|
[]
|
|
"""
|
|
n = self.start
|
|
stop = self.stop
|
|
step = self.step
|
|
cmp_ = op.lt if step > 0 else op.gt
|
|
while cmp_(n, stop):
|
|
yield n
|
|
n += step
|
|
|
|
_ops = (
|
|
(op.gt, op.ge),
|
|
(op.le, op.lt),
|
|
)
|
|
|
|
def __contains__(self, other, _ops=_ops):
|
|
# Algorithm taken from CPython
|
|
# Objects/rangeobject.c:range_contains_long
|
|
start = self.start
|
|
step = self.step
|
|
cmp_start, cmp_stop = _ops[step > 0]
|
|
return (
|
|
cmp_start(start, other) and
|
|
cmp_stop(other, self.stop) and
|
|
(other - start) % step == 0
|
|
)
|
|
|
|
del _ops
|
|
|
|
def __len__(self):
|
|
"""
|
|
Examples
|
|
--------
|
|
>>> len(range(1))
|
|
1
|
|
>>> len(range(5))
|
|
5
|
|
>>> len(range(1, 5))
|
|
4
|
|
>>> len(range(0, 5, 2))
|
|
3
|
|
>>> len(range(5, 0, -1))
|
|
5
|
|
>>> len(range(5, 0, 1))
|
|
0
|
|
"""
|
|
# Algorithm taken from CPython
|
|
# rangeobject.c:compute_range_length
|
|
step = self.step
|
|
|
|
if step > 0:
|
|
low = self.start
|
|
high = self.stop
|
|
else:
|
|
low = self.stop
|
|
high = self.start
|
|
step = -step
|
|
|
|
if low >= high:
|
|
return 0
|
|
|
|
return (high - low - 1) // step + 1
|
|
|
|
def __repr__(self):
|
|
return '%s(%s, %s%s)' % (
|
|
type(self).__name__,
|
|
self.start,
|
|
self.stop,
|
|
(', ' + str(self.step)) if self.step != 1 else '',
|
|
)
|
|
|
|
def __hash__(self):
|
|
return hash((type(self), self.start, self.stop, self.step))
|
|
|
|
def __eq__(self, other):
|
|
"""
|
|
Examples
|
|
--------
|
|
>>> range(1) == range(1)
|
|
True
|
|
>>> range(0, 5, 2) == range(0, 5, 2)
|
|
True
|
|
>>> range(5, 0, -2) == range(5, 0, -2)
|
|
True
|
|
|
|
>>> range(1) == range(2)
|
|
False
|
|
>>> range(0, 5, 2) == range(0, 5, 3)
|
|
False
|
|
"""
|
|
return all(
|
|
getattr(self, attr) == getattr(other, attr)
|
|
for attr in self.__slots__
|
|
)
|
|
else:
|
|
range = range
|
|
|
|
|
|
def from_tuple(tup):
|
|
"""Convert a tuple into a range with error handling.
|
|
|
|
Parameters
|
|
----------
|
|
tup : tuple (len 2 or 3)
|
|
The tuple to turn into a range.
|
|
|
|
Returns
|
|
-------
|
|
range : range
|
|
The range from the tuple.
|
|
|
|
Raises
|
|
------
|
|
ValueError
|
|
Raised when the tuple length is not 2 or 3.
|
|
"""
|
|
if len(tup) not in (2, 3):
|
|
raise ValueError(
|
|
'tuple must contain 2 or 3 elements, not: %d (%r' % (
|
|
len(tup),
|
|
tup,
|
|
),
|
|
)
|
|
return range(*tup)
|
|
|
|
|
|
def maybe_from_tuple(tup_or_range):
|
|
"""Convert a tuple into a range but pass ranges through silently.
|
|
|
|
This is useful to ensure that input is a range so that attributes may
|
|
be accessed with `.start`, `.stop` or so that containment checks are
|
|
constant time.
|
|
|
|
Parameters
|
|
----------
|
|
tup_or_range : tuple or range
|
|
A tuple to pass to from_tuple or a range to return.
|
|
|
|
Returns
|
|
-------
|
|
range : range
|
|
The input to convert to a range.
|
|
|
|
Raises
|
|
------
|
|
ValueError
|
|
Raised when the input is not a tuple or a range. ValueError is also
|
|
raised if the input is a tuple whose length is not 2 or 3.
|
|
"""
|
|
if isinstance(tup_or_range, tuple):
|
|
return from_tuple(tup_or_range)
|
|
elif isinstance(tup_or_range, range):
|
|
return tup_or_range
|
|
|
|
raise ValueError(
|
|
'maybe_from_tuple expects a tuple or range, got %r: %r' % (
|
|
type(tup_or_range).__name__,
|
|
tup_or_range,
|
|
),
|
|
)
|
|
|
|
|
|
def _check_steps(a, b):
|
|
"""Check that the steps of ``a`` and ``b`` are both 1.
|
|
|
|
Parameters
|
|
----------
|
|
a : range
|
|
The first range to check.
|
|
b : range
|
|
The second range to check.
|
|
|
|
Raises
|
|
------
|
|
ValueError
|
|
Raised when either step is not 1.
|
|
"""
|
|
if a.step != 1:
|
|
raise ValueError('a.step must be equal to 1, got: %s' % a.step)
|
|
if b.step != 1:
|
|
raise ValueError('b.step must be equal to 1, got: %s' % b.step)
|
|
|
|
|
|
def overlap(a, b):
|
|
"""Check if two ranges overlap.
|
|
|
|
Parameters
|
|
----------
|
|
a : range
|
|
The first range.
|
|
b : range
|
|
The second range.
|
|
|
|
Returns
|
|
-------
|
|
overlaps : bool
|
|
Do these ranges overlap.
|
|
|
|
Notes
|
|
-----
|
|
This function does not support ranges with step != 1.
|
|
"""
|
|
_check_steps(a, b)
|
|
return a.stop >= b.start and b.stop >= a.start
|
|
|
|
|
|
def merge(a, b):
|
|
"""Merge two ranges with step == 1.
|
|
|
|
Parameters
|
|
----------
|
|
a : range
|
|
The first range.
|
|
b : range
|
|
The second range.
|
|
"""
|
|
_check_steps(a, b)
|
|
return range(min(a.start, b.start), max(a.stop, b.stop))
|
|
|
|
|
|
def _combine(n, rs):
|
|
"""helper for ``_group_ranges``
|
|
"""
|
|
try:
|
|
r, rs = peek(rs)
|
|
except StopIteration:
|
|
yield n
|
|
return
|
|
|
|
if overlap(n, r):
|
|
yield merge(n, r)
|
|
next(rs)
|
|
for r in rs:
|
|
yield r
|
|
else:
|
|
yield n
|
|
for r in rs:
|
|
yield r
|
|
|
|
|
|
def group_ranges(ranges):
|
|
"""Group any overlapping ranges into a single range.
|
|
|
|
Parameters
|
|
----------
|
|
ranges : iterable[ranges]
|
|
A sorted sequence of ranges to group.
|
|
|
|
Returns
|
|
-------
|
|
grouped : iterable[ranges]
|
|
A sorted sequence of ranges with overlapping ranges merged together.
|
|
"""
|
|
return foldr(_combine, ranges, ())
|
|
|
|
|
|
def sorted_diff(rs, ss):
|
|
try:
|
|
r, rs = peek(rs)
|
|
except StopIteration:
|
|
return
|
|
|
|
try:
|
|
s, ss = peek(ss)
|
|
except StopIteration:
|
|
for r in rs:
|
|
yield r
|
|
return
|
|
|
|
rtup = (r.start, r.stop)
|
|
stup = (s.start, s.stop)
|
|
if rtup == stup:
|
|
next(rs)
|
|
next(ss)
|
|
elif rtup < stup:
|
|
yield next(rs)
|
|
else:
|
|
next(ss)
|
|
|
|
for t in sorted_diff(rs, ss):
|
|
yield t
|
|
|
|
|
|
def intersecting_ranges(ranges):
|
|
"""Return any ranges that intersect.
|
|
|
|
Parameters
|
|
----------
|
|
ranges : iterable[ranges]
|
|
A sequence of ranges to check for intersections.
|
|
|
|
Returns
|
|
-------
|
|
intersections : iterable[ranges]
|
|
A sequence of all of the ranges that intersected in ``ranges``.
|
|
|
|
Examples
|
|
--------
|
|
>>> ranges = [range(0, 1), range(2, 5), range(4, 7)]
|
|
>>> list(intersecting_ranges(ranges))
|
|
[range(2, 5), range(4, 7)]
|
|
|
|
>>> ranges = [range(0, 1), range(2, 3)]
|
|
>>> list(intersecting_ranges(ranges))
|
|
[]
|
|
|
|
>>> ranges = [range(0, 1), range(1, 2)]
|
|
>>> list(intersecting_ranges(ranges))
|
|
[range(0, 1), range(1, 2)]
|
|
"""
|
|
ranges = sorted(ranges, key=op.attrgetter('start'))
|
|
return sorted_diff(ranges, group_ranges(ranges))
|