35
35
from pandas .util ._decorators import cache_readonly , Appender
36
36
from pandas .core .config import get_option
37
37
from pandas .tseries .frequencies import to_offset
38
+ from pandas .tseries .offsets import DateOffset
38
39
39
40
import pandas .core .indexes .base as ibase
40
41
_index_doc_kwargs = dict (ibase ._index_doc_kwargs )
@@ -1033,6 +1034,24 @@ def func(self, other):
1033
1034
IntervalIndex ._add_logical_methods_disabled ()
1034
1035
1035
1036
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
+
1036
1055
def interval_range (start = None , end = None , periods = None , freq = None ,
1037
1056
name = None , closed = 'right' ):
1038
1057
"""
@@ -1107,27 +1126,15 @@ def interval_range(start=None, end=None, periods=None, freq=None,
1107
1126
raise ValueError ('Of the three parameters: start, end, and periods, '
1108
1127
'exactly two must be specified' )
1109
1128
1110
- iv_type = {'numeric' : True , 'timestamp' : True , 'timedelta' : True }
1111
-
1112
1129
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 ):
1120
1134
msg = 'start must be numeric or datetime-like, got {start}'
1121
1135
raise ValueError (msg .format (start = start ))
1122
1136
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 ):
1131
1138
msg = 'end must be numeric or datetime-like, got {end}'
1132
1139
raise ValueError (msg .format (end = end ))
1133
1140
@@ -1137,22 +1144,21 @@ def interval_range(start=None, end=None, periods=None, freq=None,
1137
1144
msg = 'periods must be a number, got {periods}'
1138
1145
raise TypeError (msg .format (periods = periods ))
1139
1146
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 ):
1144
1149
try :
1145
1150
freq = to_offset (freq )
1146
- iv_type ['numeric' ] = False
1147
1151
except ValueError :
1148
1152
raise ValueError ('freq must be numeric or convertible to '
1149
1153
'DateOffset, got {freq}' .format (freq = freq ))
1150
1154
1151
1155
# 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 )]):
1153
1159
raise TypeError ("start, end, freq need to be type compatible" )
1154
1160
1155
- if iv_type [ 'numeric' ] :
1161
+ if is_number ( endpoint ) :
1156
1162
if periods is None :
1157
1163
periods = int ((end - start ) // freq )
1158
1164
@@ -1164,7 +1170,7 @@ def interval_range(start=None, end=None, periods=None, freq=None,
1164
1170
1165
1171
# end + freq for inclusive endpoint
1166
1172
breaks = np .arange (start , end + freq , freq )
1167
- elif iv_type [ 'timestamp' ] :
1173
+ elif isinstance ( endpoint , Timestamp ) :
1168
1174
# add one to account for interval endpoints (n breaks = n-1 intervals)
1169
1175
if periods is not None :
1170
1176
periods += 1
0 commit comments