Skip to content

Commit d5579cb

Browse files
committed
clean interval_range
1 parent 549c402 commit d5579cb

File tree

1 file changed

+31
-25
lines changed

1 file changed

+31
-25
lines changed

pandas/core/indexes/interval.py

+31-25
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
from pandas.util._decorators import cache_readonly, Appender
3636
from pandas.core.config import get_option
3737
from pandas.tseries.frequencies import to_offset
38+
from pandas.tseries.offsets import DateOffset
3839

3940
import pandas.core.indexes.base as ibase
4041
_index_doc_kwargs = dict(ibase._index_doc_kwargs)
@@ -1033,6 +1034,24 @@ def func(self, other):
10331034
IntervalIndex._add_logical_methods_disabled()
10341035

10351036

1037+
def _is_valid_endpoint(endpoint):
1038+
"""helper for interval_range to check if start/end are valid types"""
1039+
return any([is_number(endpoint),
1040+
isinstance(endpoint, Timestamp),
1041+
isinstance(endpoint, Timedelta),
1042+
endpoint is None])
1043+
1044+
1045+
def _is_type_compatible(a, b):
1046+
"""helper for interval_range to check type compat of start/end/freq"""
1047+
is_ts_compat = lambda x: isinstance(x, (Timestamp, DateOffset))
1048+
is_td_compat = lambda x: isinstance(x, (Timedelta, DateOffset))
1049+
return ((is_number(a) and is_number(b)) or
1050+
(is_ts_compat(a) and is_ts_compat(b)) or
1051+
(is_td_compat(a) and is_td_compat(b)) or
1052+
com._any_none(a, b))
1053+
1054+
10361055
def interval_range(start=None, end=None, periods=None, freq=None,
10371056
name=None, closed='right'):
10381057
"""
@@ -1107,27 +1126,15 @@ def interval_range(start=None, end=None, periods=None, freq=None,
11071126
raise ValueError('Of the three parameters: start, end, and periods, '
11081127
'exactly two must be specified')
11091128

1110-
iv_type = {'numeric': True, 'timestamp': True, 'timedelta': True}
1111-
11121129
start = com._maybe_box_datetimelike(start)
1113-
if is_number(start):
1114-
iv_type.update({k: False for k in iv_type if k != 'numeric'})
1115-
elif isinstance(start, Timestamp):
1116-
iv_type.update({k: False for k in iv_type if k != 'timestamp'})
1117-
elif isinstance(start, Timedelta):
1118-
iv_type.update({k: False for k in iv_type if k != 'timedelta'})
1119-
elif start is not None:
1130+
end = com._maybe_box_datetimelike(end)
1131+
endpoint = next(com._not_none(start, end))
1132+
1133+
if not _is_valid_endpoint(start):
11201134
msg = 'start must be numeric or datetime-like, got {start}'
11211135
raise ValueError(msg.format(start=start))
11221136

1123-
end = com._maybe_box_datetimelike(end)
1124-
if is_number(end):
1125-
iv_type.update({k: False for k in iv_type if k != 'numeric'})
1126-
elif isinstance(end, Timestamp):
1127-
iv_type.update({k: False for k in iv_type if k != 'timestamp'})
1128-
elif isinstance(end, Timedelta):
1129-
iv_type.update({k: False for k in iv_type if k != 'timedelta'})
1130-
elif end is not None:
1137+
if not _is_valid_endpoint(end):
11311138
msg = 'end must be numeric or datetime-like, got {end}'
11321139
raise ValueError(msg.format(end=end))
11331140

@@ -1137,22 +1144,21 @@ def interval_range(start=None, end=None, periods=None, freq=None,
11371144
msg = 'periods must be a number, got {periods}'
11381145
raise TypeError(msg.format(periods=periods))
11391146

1140-
freq = freq or (1 if iv_type['numeric'] else 'D')
1141-
if is_number(freq):
1142-
iv_type.update({k: False for k in iv_type if k != 'numeric'})
1143-
else:
1147+
freq = freq or (1 if is_number(endpoint) else 'D')
1148+
if not is_number(freq):
11441149
try:
11451150
freq = to_offset(freq)
1146-
iv_type['numeric'] = False
11471151
except ValueError:
11481152
raise ValueError('freq must be numeric or convertible to '
11491153
'DateOffset, got {freq}'.format(freq=freq))
11501154

11511155
# verify type compatibility
1152-
if not any(iv_type.values()):
1156+
if not all([_is_type_compatible(start, end),
1157+
_is_type_compatible(start, freq),
1158+
_is_type_compatible(end, freq)]):
11531159
raise TypeError("start, end, freq need to be type compatible")
11541160

1155-
if iv_type['numeric']:
1161+
if is_number(endpoint):
11561162
if periods is None:
11571163
periods = int((end - start) // freq)
11581164

@@ -1164,7 +1170,7 @@ def interval_range(start=None, end=None, periods=None, freq=None,
11641170

11651171
# end + freq for inclusive endpoint
11661172
breaks = np.arange(start, end + freq, freq)
1167-
elif iv_type['timestamp']:
1173+
elif isinstance(endpoint, Timestamp):
11681174
# add one to account for interval endpoints (n breaks = n-1 intervals)
11691175
if periods is not None:
11701176
periods += 1

0 commit comments

Comments
 (0)