Index tests indexing (pandas-dev#18145)

Author: jbrockmendel

pandas/tests/indexes/datetimes/test_datetime.py

@@ -410,89 +410,6 @@ def test_sort_values(self):
         assert ordered[::-1].is_monotonic
         tm.assert_numpy_array_equal(dexer, np.array([0, 2, 1], dtype=np.intp))
 
-    def test_take(self):
-        dates = [datetime(2010, 1, 1, 14), datetime(2010, 1, 1, 15),
-                 datetime(2010, 1, 1, 17), datetime(2010, 1, 1, 21)]
-
-        for tz in [None, 'US/Eastern', 'Asia/Tokyo']:
-            idx = DatetimeIndex(start='2010-01-01 09:00',
-                                end='2010-02-01 09:00', freq='H', tz=tz,
-                                name='idx')
-            expected = DatetimeIndex(dates, freq=None, name='idx', tz=tz)
-
-            taken1 = idx.take([5, 6, 8, 12])
-            taken2 = idx[[5, 6, 8, 12]]
-
-            for taken in [taken1, taken2]:
-                tm.assert_index_equal(taken, expected)
-                assert isinstance(taken, DatetimeIndex)
-                assert taken.freq is None
-                assert taken.tz == expected.tz
-                assert taken.name == expected.name
-
-    def test_take_fill_value(self):
-        # GH 12631
-        idx = pd.DatetimeIndex(['2011-01-01', '2011-02-01', '2011-03-01'],
-                               name='xxx')
-        result = idx.take(np.array([1, 0, -1]))
-        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'],
-                                    name='xxx')
-        tm.assert_index_equal(result, expected)
-
-        # fill_value
-        result = idx.take(np.array([1, 0, -1]), fill_value=True)
-        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', 'NaT'],
-                                    name='xxx')
-        tm.assert_index_equal(result, expected)
-
-        # allow_fill=False
-        result = idx.take(np.array([1, 0, -1]), allow_fill=False,
-                          fill_value=True)
-        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'],
-                                    name='xxx')
-        tm.assert_index_equal(result, expected)
-
-        msg = ('When allow_fill=True and fill_value is not None, '
-               'all indices must be >= -1')
-        with tm.assert_raises_regex(ValueError, msg):
-            idx.take(np.array([1, 0, -2]), fill_value=True)
-        with tm.assert_raises_regex(ValueError, msg):
-            idx.take(np.array([1, 0, -5]), fill_value=True)
-
-        with pytest.raises(IndexError):
-            idx.take(np.array([1, -5]))
-
-    def test_take_fill_value_with_timezone(self):
-        idx = pd.DatetimeIndex(['2011-01-01', '2011-02-01', '2011-03-01'],
-                               name='xxx', tz='US/Eastern')
-        result = idx.take(np.array([1, 0, -1]))
-        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'],
-                                    name='xxx', tz='US/Eastern')
-        tm.assert_index_equal(result, expected)
-
-        # fill_value
-        result = idx.take(np.array([1, 0, -1]), fill_value=True)
-        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', 'NaT'],
-                                    name='xxx', tz='US/Eastern')
-        tm.assert_index_equal(result, expected)
-
-        # allow_fill=False
-        result = idx.take(np.array([1, 0, -1]), allow_fill=False,
-                          fill_value=True)
-        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'],
-                                    name='xxx', tz='US/Eastern')
-        tm.assert_index_equal(result, expected)
-
-        msg = ('When allow_fill=True and fill_value is not None, '
-               'all indices must be >= -1')
-        with tm.assert_raises_regex(ValueError, msg):
-            idx.take(np.array([1, 0, -2]), fill_value=True)
-        with tm.assert_raises_regex(ValueError, msg):
-            idx.take(np.array([1, 0, -5]), fill_value=True)
-
-        with pytest.raises(IndexError):
-            idx.take(np.array([1, -5]))
-
     def test_map_bug_1677(self):
         index = DatetimeIndex(['2012-04-25 09:30:00.393000'])
         f = index.asof

pandas/tests/indexes/datetimes/test_indexing.py

@@ -1,11 +1,15 @@
+from datetime import datetime
 import pytest
 
 import pytz
 import numpy as np
 import pandas as pd
 import pandas.util.testing as tm
 import pandas.compat as compat
-from pandas import notna, Index, DatetimeIndex, datetime, date_range
+from pandas import notna, Index, DatetimeIndex, date_range, Timestamp
+from pandas.tseries.offsets import CDay, BDay
+
+START, END = datetime(2009, 1, 1), datetime(2010, 1, 1)
 
 
 class TestDatetimeIndex(object):
@@ -239,3 +243,234 @@ def test_delete_slice(self):
             assert result.name == expected.name
             assert result.freq == expected.freq
             assert result.tz == expected.tz
+
+    def test_getitem(self):
+        idx1 = pd.date_range('2011-01-01', '2011-01-31', freq='D', name='idx')
+        idx2 = pd.date_range('2011-01-01', '2011-01-31', freq='D',
+                             tz='Asia/Tokyo', name='idx')
+
+        for idx in [idx1, idx2]:
+            result = idx[0]
+            assert result == Timestamp('2011-01-01', tz=idx.tz)
+
+            result = idx[0:5]
+            expected = pd.date_range('2011-01-01', '2011-01-05', freq='D',
+                                     tz=idx.tz, name='idx')
+            tm.assert_index_equal(result, expected)
+            assert result.freq == expected.freq
+
+            result = idx[0:10:2]
+            expected = pd.date_range('2011-01-01', '2011-01-09', freq='2D',
+                                     tz=idx.tz, name='idx')
+            tm.assert_index_equal(result, expected)
+            assert result.freq == expected.freq
+
+            result = idx[-20:-5:3]
+            expected = pd.date_range('2011-01-12', '2011-01-24', freq='3D',
+                                     tz=idx.tz, name='idx')
+            tm.assert_index_equal(result, expected)
+            assert result.freq == expected.freq
+
+            result = idx[4::-1]
+            expected = DatetimeIndex(['2011-01-05', '2011-01-04', '2011-01-03',
+                                      '2011-01-02', '2011-01-01'],
+                                     freq='-1D', tz=idx.tz, name='idx')
+            tm.assert_index_equal(result, expected)
+            assert result.freq == expected.freq
+
+    def test_take(self):
+        # GH 10295
+        idx1 = pd.date_range('2011-01-01', '2011-01-31', freq='D', name='idx')
+        idx2 = pd.date_range('2011-01-01', '2011-01-31', freq='D',
+                             tz='Asia/Tokyo', name='idx')
+
+        for idx in [idx1, idx2]:
+            result = idx.take([0])
+            assert result == Timestamp('2011-01-01', tz=idx.tz)
+
+            result = idx.take([0, 1, 2])
+            expected = pd.date_range('2011-01-01', '2011-01-03', freq='D',
+                                     tz=idx.tz, name='idx')
+            tm.assert_index_equal(result, expected)
+            assert result.freq == expected.freq
+
+            result = idx.take([0, 2, 4])
+            expected = pd.date_range('2011-01-01', '2011-01-05', freq='2D',
+                                     tz=idx.tz, name='idx')
+            tm.assert_index_equal(result, expected)
+            assert result.freq == expected.freq
+
+            result = idx.take([7, 4, 1])
+            expected = pd.date_range('2011-01-08', '2011-01-02', freq='-3D',
+                                     tz=idx.tz, name='idx')
+            tm.assert_index_equal(result, expected)
+            assert result.freq == expected.freq
+
+            result = idx.take([3, 2, 5])
+            expected = DatetimeIndex(['2011-01-04', '2011-01-03',
+                                      '2011-01-06'],
+                                     freq=None, tz=idx.tz, name='idx')
+            tm.assert_index_equal(result, expected)
+            assert result.freq is None
+
+            result = idx.take([-3, 2, 5])
+            expected = DatetimeIndex(['2011-01-29', '2011-01-03',
+                                      '2011-01-06'],
+                                     freq=None, tz=idx.tz, name='idx')
+            tm.assert_index_equal(result, expected)
+            assert result.freq is None
+
+    def test_take_invalid_kwargs(self):
+        idx = pd.date_range('2011-01-01', '2011-01-31', freq='D', name='idx')
+        indices = [1, 6, 5, 9, 10, 13, 15, 3]
+
+        msg = r"take\(\) got an unexpected keyword argument 'foo'"
+        tm.assert_raises_regex(TypeError, msg, idx.take,
+                               indices, foo=2)
+
+        msg = "the 'out' parameter is not supported"
+        tm.assert_raises_regex(ValueError, msg, idx.take,
+                               indices, out=indices)
+
+        msg = "the 'mode' parameter is not supported"
+        tm.assert_raises_regex(ValueError, msg, idx.take,
+                               indices, mode='clip')
+
+    # TODO: This method came from test_datetime; de-dup with version above
+    def test_take2(self):
+        dates = [datetime(2010, 1, 1, 14), datetime(2010, 1, 1, 15),
+                 datetime(2010, 1, 1, 17), datetime(2010, 1, 1, 21)]
+
+        for tz in [None, 'US/Eastern', 'Asia/Tokyo']:
+            idx = DatetimeIndex(start='2010-01-01 09:00',
+                                end='2010-02-01 09:00', freq='H', tz=tz,
+                                name='idx')
+            expected = DatetimeIndex(dates, freq=None, name='idx', tz=tz)
+
+            taken1 = idx.take([5, 6, 8, 12])
+            taken2 = idx[[5, 6, 8, 12]]
+
+            for taken in [taken1, taken2]:
+                tm.assert_index_equal(taken, expected)
+                assert isinstance(taken, DatetimeIndex)
+                assert taken.freq is None
+                assert taken.tz == expected.tz
+                assert taken.name == expected.name
+
+    def test_take_fill_value(self):
+        # GH 12631
+        idx = pd.DatetimeIndex(['2011-01-01', '2011-02-01', '2011-03-01'],
+                               name='xxx')
+        result = idx.take(np.array([1, 0, -1]))
+        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'],
+                                    name='xxx')
+        tm.assert_index_equal(result, expected)
+
+        # fill_value
+        result = idx.take(np.array([1, 0, -1]), fill_value=True)
+        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', 'NaT'],
+                                    name='xxx')
+        tm.assert_index_equal(result, expected)
+
+        # allow_fill=False
+        result = idx.take(np.array([1, 0, -1]), allow_fill=False,
+                          fill_value=True)
+        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'],
+                                    name='xxx')
+        tm.assert_index_equal(result, expected)
+
+        msg = ('When allow_fill=True and fill_value is not None, '
+               'all indices must be >= -1')
+        with tm.assert_raises_regex(ValueError, msg):
+            idx.take(np.array([1, 0, -2]), fill_value=True)
+        with tm.assert_raises_regex(ValueError, msg):
+            idx.take(np.array([1, 0, -5]), fill_value=True)
+
+        with pytest.raises(IndexError):
+            idx.take(np.array([1, -5]))
+
+    def test_take_fill_value_with_timezone(self):
+        idx = pd.DatetimeIndex(['2011-01-01', '2011-02-01', '2011-03-01'],
+                               name='xxx', tz='US/Eastern')
+        result = idx.take(np.array([1, 0, -1]))
+        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'],
+                                    name='xxx', tz='US/Eastern')
+        tm.assert_index_equal(result, expected)
+
+        # fill_value
+        result = idx.take(np.array([1, 0, -1]), fill_value=True)
+        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', 'NaT'],
+                                    name='xxx', tz='US/Eastern')
+        tm.assert_index_equal(result, expected)
+
+        # allow_fill=False
+        result = idx.take(np.array([1, 0, -1]), allow_fill=False,
+                          fill_value=True)
+        expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'],
+                                    name='xxx', tz='US/Eastern')
+        tm.assert_index_equal(result, expected)
+
+        msg = ('When allow_fill=True and fill_value is not None, '
+               'all indices must be >= -1')
+        with tm.assert_raises_regex(ValueError, msg):
+            idx.take(np.array([1, 0, -2]), fill_value=True)
+        with tm.assert_raises_regex(ValueError, msg):
+            idx.take(np.array([1, 0, -5]), fill_value=True)
+
+        with pytest.raises(IndexError):
+            idx.take(np.array([1, -5]))
+
+
+class TestBusinessDatetimeIndexIndexing(object):
+    def setup_method(self, method):
+        self.rng = pd.bdate_range(START, END)
+
+    def test_getitem(self):
+        smaller = self.rng[:5]
+        exp = DatetimeIndex(self.rng.view(np.ndarray)[:5])
+        tm.assert_index_equal(smaller, exp)
+
+        assert smaller.offset == self.rng.offset
+
+        sliced = self.rng[::5]
+        assert sliced.offset == BDay() * 5
+
+        fancy_indexed = self.rng[[4, 3, 2, 1, 0]]
+        assert len(fancy_indexed) == 5
+        assert isinstance(fancy_indexed, DatetimeIndex)
+        assert fancy_indexed.freq is None
+
+        # 32-bit vs. 64-bit platforms
+        assert self.rng[4] == self.rng[np.int_(4)]
+
+    def test_getitem_matplotlib_hackaround(self):
+        values = self.rng[:, None]
+        expected = self.rng.values[:, None]
+        tm.assert_numpy_array_equal(values, expected)
+
+
+class TestCustomDatetimeIndexIndexing(object):
+    def setup_method(self, method):
+        self.rng = pd.bdate_range(START, END, freq='C')
+
+    def test_getitem(self):
+        smaller = self.rng[:5]
+        exp = DatetimeIndex(self.rng.view(np.ndarray)[:5])
+        tm.assert_index_equal(smaller, exp)
+        assert smaller.offset == self.rng.offset
+
+        sliced = self.rng[::5]
+        assert sliced.offset == CDay() * 5
+
+        fancy_indexed = self.rng[[4, 3, 2, 1, 0]]
+        assert len(fancy_indexed) == 5
+        assert isinstance(fancy_indexed, DatetimeIndex)
+        assert fancy_indexed.freq is None
+
+        # 32-bit vs. 64-bit platforms
+        assert self.rng[4] == self.rng[np.int_(4)]
+
+    def test_getitem_matplotlib_hackaround(self):
+        values = self.rng[:, None]
+        expected = self.rng.values[:, None]
+        tm.assert_numpy_array_equal(values, expected)