fix datetimelike

Author: jschendel

pandas/core/indexes/interval.py

@@ -28,12 +28,12 @@
     intervals_to_interval_bounds)
 
 from pandas.core.indexes.datetimes import date_range
+from pandas.core.indexes.timedeltas import timedelta_range
 from pandas.core.indexes.multi import MultiIndex
 from pandas.compat.numpy import function as nv
 from pandas.core import common as com
 from pandas.util._decorators import cache_readonly, Appender
 from pandas.core.config import get_option
-from pandas.tseries.offsets import DateOffset
 from pandas.tseries.frequencies import to_offset
 
 import pandas.core.indexes.base as ibase
@@ -1107,48 +1107,52 @@ def interval_range(start=None, end=None, periods=None, freq=None,
         raise ValueError('Of the three parameters: start, end, and periods, '
                          'exactly two must be specified')
 
-    # assume datetime-like unless we find numeric start or end
-    is_datetime_interval = True
+    iv_type = {'numeric': True, 'timestamp': True, 'timedelta': True}
 
+    start = com._maybe_box_datetimelike(start)
     if is_number(start):
-        is_datetime_interval = False
+        iv_type.update({k: False for k in iv_type if k != 'numeric'})
+    elif isinstance(start, Timestamp):
+        iv_type.update({k: False for k in iv_type if k != 'timestamp'})
+    elif isinstance(start, Timedelta):
+        iv_type.update({k: False for k in iv_type if k != 'timedelta'})
     elif start is not None:
-        try:
-            start = Timestamp(start)
-        except (TypeError, ValueError):
-            raise ValueError('start must be numeric or datetime-like')
+        msg = 'start must be numeric or datetime-like, got {start}'
+        raise ValueError(msg.format(start=start))
 
+    end = com._maybe_box_datetimelike(end)
     if is_number(end):
-        is_datetime_interval = False
+        iv_type.update({k: False for k in iv_type if k != 'numeric'})
+    elif isinstance(end, Timestamp):
+        iv_type.update({k: False for k in iv_type if k != 'timestamp'})
+    elif isinstance(end, Timedelta):
+        iv_type.update({k: False for k in iv_type if k != 'timedelta'})
     elif end is not None:
-        try:
-            end = Timestamp(end)
-        except (TypeError, ValueError):
-            raise ValueError('end must be numeric or datetime-like')
+        msg = 'end must be numeric or datetime-like, got {end}'
+        raise ValueError(msg.format(end=end))
 
     if is_float(periods):
         periods = int(periods)
     elif not is_integer(periods) and periods is not None:
         msg = 'periods must be a number, got {periods}'
         raise TypeError(msg.format(periods=periods))
 
-    if is_datetime_interval:
-        freq = freq or 'D'
-        if not isinstance(freq, DateOffset):
-            try:
-                freq = to_offset(freq)
-            except ValueError:
-                raise ValueError('freq must be convertible to DateOffset when '
-                                 'start/end are datetime-like')
+    freq = freq or (1 if iv_type['numeric'] else 'D')
+    if is_number(freq):
+        iv_type.update({k: False for k in iv_type if k != 'numeric'})
     else:
-        freq = freq or 1
+        try:
+            freq = to_offset(freq)
+            iv_type['numeric'] = False
+        except ValueError:
+            raise ValueError('freq must be numeric or convertible to '
+                             'DateOffset, got {freq}'.format(freq=freq))
 
     # verify type compatibility
-    is_numeric_interval = all(map(is_number, com._not_none(start, end, freq)))
-    if not is_datetime_interval and not is_numeric_interval:
+    if not any(iv_type.values()):
         raise TypeError("start, end, freq need to be type compatible")
 
-    if is_numeric_interval:
+    if iv_type['numeric']:
         if periods is None:
             periods = int((end - start) // freq)
 
@@ -1160,10 +1164,16 @@ def interval_range(start=None, end=None, periods=None, freq=None,
 
         # end + freq for inclusive endpoint
         breaks = np.arange(start, end + freq, freq)
-    else:
+    elif iv_type['timestamp']:
         # add one to account for interval endpoints (n breaks = n-1 intervals)
         if periods is not None:
             periods += 1
         breaks = date_range(start=start, end=end, periods=periods, freq=freq)
+    else:
+        # add one to account for interval endpoints (n breaks = n-1 intervals)
+        if periods is not None:
+            periods += 1
+        breaks = timedelta_range(start=start, end=end, periods=periods,
+                                 freq=freq)
 
     return IntervalIndex.from_breaks(breaks, name=name, closed=closed)

pandas/tests/indexes/test_interval.py

@@ -2,10 +2,11 @@
 
 import pytest
 import numpy as np
-import datetime
+from datetime import timedelta
 from pandas import (Interval, IntervalIndex, Index, isna,
                     interval_range, Timestamp, Timedelta,
-                    compat, date_range, DateOffset)
+                    compat, date_range, timedelta_range, DateOffset)
+from pandas.tseries.offsets import Day
 from pandas._libs.interval import IntervalTree
 from pandas.tests.indexes.common import Base
 import pandas.util.testing as tm
@@ -763,8 +764,8 @@ def test_construction_from_numeric(self, closed):
     def test_construction_from_timestamp(self, closed):
         # combinations of start/end/periods without freq
         start, end = Timestamp('2017-01-01'), Timestamp('2017-01-06')
-        expected = IntervalIndex.from_breaks(date_range(start=start, end=end),
-                                             name='foo', closed=closed)
+        breaks = date_range(start=start, end=end)
+        expected = IntervalIndex.from_breaks(breaks, name='foo', closed=closed)
 
         result = interval_range(start=start, end=end, name='foo',
                                 closed=closed)
@@ -826,37 +827,90 @@ def test_construction_from_timestamp(self, closed):
                                 closed=closed)
         tm.assert_index_equal(result, expected)
 
+    @pytest.mark.parametrize('closed', ['left', 'right', 'neither', 'both'])
+    def test_construction_from_timedelta(self, closed):
+        # combinations of start/end/periods without freq
+        start, end = Timedelta('1 day'), Timedelta('6 days')
+        breaks = timedelta_range(start=start, end=end)
+        expected = IntervalIndex.from_breaks(breaks, name='foo', closed=closed)
+
+        result = interval_range(start=start, end=end, name='foo',
+                                closed=closed)
+        tm.assert_index_equal(result, expected)
+
+        result = interval_range(start=start, periods=5, name='foo',
+                                closed=closed)
+        tm.assert_index_equal(result, expected)
+
+        result = interval_range(end=end, periods=5, name='foo',
+                                closed=closed)
+        tm.assert_index_equal(result, expected)
+
+        # combinations of start/end/periods with fixed freq
+        freq = '2D'
+        start, end = Timedelta('1 day'), Timedelta('7 days')
+        breaks = timedelta_range(start=start, end=end, freq=freq)
+        expected = IntervalIndex.from_breaks(breaks, name='foo', closed=closed)
+
+        result = interval_range(start=start, end=end, freq=freq, name='foo',
+                                closed=closed)
+        tm.assert_index_equal(result, expected)
+
+        result = interval_range(start=start, periods=3, freq=freq, name='foo',
+                                closed=closed)
+        tm.assert_index_equal(result, expected)
+
+        result = interval_range(end=end, periods=3, freq=freq, name='foo',
+                                closed=closed)
+        tm.assert_index_equal(result, expected)
+
+        # output truncates early if freq causes end to be skipped.
+        end = Timedelta('7 days 1 hour')
+        result = interval_range(start=start, end=end, freq=freq, name='foo',
+                                closed=closed)
+        tm.assert_index_equal(result, expected)
+
     def test_constructor_coverage(self):
         # float value for periods
         expected = pd.interval_range(start=0, periods=10)
         result = pd.interval_range(start=0, periods=10.5)
         tm.assert_index_equal(result, expected)
 
-        # equivalent datetime-like start/end
+        # equivalent timestamp-like start/end
         start, end = Timestamp('2017-01-01'), Timestamp('2017-01-15')
         expected = pd.interval_range(start=start, end=end)
 
         result = pd.interval_range(start=start.to_pydatetime(),
                                    end=end.to_pydatetime())
         tm.assert_index_equal(result, expected)
 
-        result = pd.interval_range(start=start.date(), end=end.date())
+        result = pd.interval_range(start=start.tz_localize('UTC'),
+                                   end=end.tz_localize('UTC'))
         tm.assert_index_equal(result, expected)
 
-        result = pd.interval_range(start=str(start), end=str(end))
+        result = pd.interval_range(start=start.asm8, end=end.asm8)
         tm.assert_index_equal(result, expected)
 
-        result = pd.interval_range(start=start.strftime('%Y-%m-%d'),
-                                   end=end.strftime('%Y-%m-%d'))
+        # equivalent freq with timestamp
+        equiv_freq = ['D', Day(), Timedelta(days=1), timedelta(days=1),
+                      DateOffset(days=1)]
+        for freq in equiv_freq:
+            result = pd.interval_range(start=start, end=end, freq=freq)
+            tm.assert_index_equal(result, expected)
+
+        # equivalent timedelta-like start/end
+        start, end = Timedelta(days=1), Timedelta(days=10)
+        expected = pd.interval_range(start=start, end=end)
+
+        result = pd.interval_range(start=start.to_pytimedelta(),
+                                   end=end.to_pytimedelta())
         tm.assert_index_equal(result, expected)
 
-        result = pd.interval_range(start=start.strftime('%m/%d/%y'),
-                                   end=end.strftime('%m/%d/%y'))
+        result = pd.interval_range(start=start.asm8, end=end.asm8)
         tm.assert_index_equal(result, expected)
 
-        # equivalent freq
-        equiv_freq = ['D', DateOffset(days=1), Timedelta(days=1),
-                      datetime.timedelta(days=1)]
+        # equivalent freq with timedelta
+        equiv_freq = ['D', Day(), Timedelta(days=1), timedelta(days=1)]
         for freq in equiv_freq:
             result = pd.interval_range(start=start, end=end, freq=freq)
             tm.assert_index_equal(result, expected)
@@ -885,37 +939,61 @@ def test_errors(self):
         # mixed units
         msg = 'start, end, freq need to be type compatible'
         with tm.assert_raises_regex(TypeError, msg):
-            interval_range(start=Timestamp('20130101'), end=10, freq=2)
+            interval_range(start=0, end=Timestamp('20130101'), freq=2)
+
+        with tm.assert_raises_regex(TypeError, msg):
+            interval_range(start=0, end=Timedelta('1 day'), freq=2)
 
         with tm.assert_raises_regex(TypeError, msg):
-            interval_range(start=0, end=Timestamp('20130101'), freq=2)
+            interval_range(start=0, end=10, freq='D')
+
+        with tm.assert_raises_regex(TypeError, msg):
+            interval_range(start=Timestamp('20130101'), end=10, freq='D')
+
+        with tm.assert_raises_regex(TypeError, msg):
+            interval_range(start=Timestamp('20130101'),
+                           end=Timedelta('1 day'), freq='D')
 
         with tm.assert_raises_regex(TypeError, msg):
-            interval_range(start=0, end=10, freq=Timedelta('1day'))
+            interval_range(start=Timestamp('20130101'),
+                           end=Timestamp('20130110'), freq=2)
+
+        with tm.assert_raises_regex(TypeError, msg):
+            interval_range(start=Timedelta('1 day'), end=10, freq='D')
+
+        with tm.assert_raises_regex(TypeError, msg):
+            interval_range(start=Timedelta('1 day'),
+                           end=Timestamp('20130110'), freq='D')
+
+        with tm.assert_raises_regex(TypeError, msg):
+            interval_range(start=Timedelta('1 day'),
+                           end=Timedelta('10 days'), freq=2)
 
         # invalid periods
         msg = 'periods must be a number, got foo'
         with tm.assert_raises_regex(TypeError, msg):
             interval_range(start=0, periods='foo')
 
         # invalid start
-        msg = 'start must be numeric or datetime-like'
+        msg = 'start must be numeric or datetime-like, got foo'
         with tm.assert_raises_regex(ValueError, msg):
             interval_range(start='foo', periods=10)
 
         # invalid end
-        msg = 'end must be numeric or datetime-like'
+        msg = 'end must be numeric or datetime-like, got \(0, 1\]'
         with tm.assert_raises_regex(ValueError, msg):
             interval_range(end=Interval(0, 1), periods=10)
 
         # invalid freq for datetime-like
-        msg = ('freq must be convertible to DateOffset when start/end are '
-               'datetime-like')
+        msg = 'freq must be numeric or convertible to DateOffset, got foo'
+        with tm.assert_raises_regex(ValueError, msg):
+            interval_range(start=0, end=10, freq='foo')
+
         with tm.assert_raises_regex(ValueError, msg):
             interval_range(start=Timestamp('20130101'), periods=10, freq='foo')
 
         with tm.assert_raises_regex(ValueError, msg):
-            interval_range(end=Timestamp('20130101'), periods=10, freq='foo')
+            interval_range(end=Timedelta('1 day'), periods=10, freq='foo')
 
 
 class TestIntervalTree(object):