TST/REF: collect Index tests by method (#33870)

Author: jbrockmendel

pandas/tests/indexes/datetimes/test_datetime.py

@@ -5,46 +5,13 @@
 import pytest
 
 import pandas as pd
-from pandas import DataFrame, DatetimeIndex, Index, NaT, Timestamp, date_range, offsets
+from pandas import DataFrame, DatetimeIndex, Index, Timestamp, date_range, offsets
 import pandas._testing as tm
 
 randn = np.random.randn
 
 
 class TestDatetimeIndex:
-    def test_roundtrip_pickle_with_tz(self):
-
-        # GH 8367
-        # round-trip of timezone
-        index = date_range("20130101", periods=3, tz="US/Eastern", name="foo")
-        unpickled = tm.round_trip_pickle(index)
-        tm.assert_index_equal(index, unpickled)
-
-    def test_pickle(self):
-
-        # GH#4606
-        p = tm.round_trip_pickle(NaT)
-        assert p is NaT
-
-        idx = pd.to_datetime(["2013-01-01", NaT, "2014-01-06"])
-        idx_p = tm.round_trip_pickle(idx)
-        assert idx_p[0] == idx[0]
-        assert idx_p[1] is NaT
-        assert idx_p[2] == idx[2]
-
-        # GH#11002
-        # don't infer freq
-        idx = date_range("1750-1-1", "2050-1-1", freq="7D")
-        idx_p = tm.round_trip_pickle(idx)
-        tm.assert_index_equal(idx, idx_p)
-
-    def test_pickle_after_set_freq(self):
-        dti = date_range("20130101", periods=3, tz="US/Eastern", name="foo")
-        dti = dti._with_freq(None)
-
-        res = tm.round_trip_pickle(dti)
-        tm.assert_index_equal(res, dti)
-
     def test_reindex_preserves_tz_if_target_is_empty_list_or_array(self):
         # GH7774
         index = date_range("20130101", periods=3, tz="US/Eastern")
@@ -164,23 +131,6 @@ def test_append_nondatetimeindex(self):
         result = rng.append(idx)
         assert isinstance(result[0], Timestamp)
 
-    def test_map(self):
-        rng = date_range("1/1/2000", periods=10)
-
-        f = lambda x: x.strftime("%Y%m%d")
-        result = rng.map(f)
-        exp = Index([f(x) for x in rng], dtype="<U8")
-        tm.assert_index_equal(result, exp)
-
-    def test_map_fallthrough(self, capsys):
-        # GH#22067, check we don't get warnings about silently ignored errors
-        dti = date_range("2017-01-01", "2018-01-01", freq="B")
-
-        dti.map(lambda x: pd.Period(year=x.year, month=x.month, freq="M"))
-
-        captured = capsys.readouterr()
-        assert captured.err == ""
-
     def test_iteration_preserves_tz(self):
         # see gh-8890
         index = date_range("2012-01-01", periods=3, freq="H", tz="US/Eastern")
@@ -264,14 +214,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_map_bug_1677(self):
-        index = DatetimeIndex(["2012-04-25 09:30:00.393000"])
-        f = index.asof
-
-        result = index.map(f)
-        expected = Index([f(index[0])])
-        tm.assert_index_equal(result, expected)
-
     def test_groupby_function_tuple_1677(self):
         df = DataFrame(np.random.rand(100), index=date_range("1/1/2000", periods=100))
         monthly_group = df.groupby(lambda x: (x.year, x.month))
@@ -454,18 +396,6 @@ def test_to_frame_datetime_tz(self):
         expected = DataFrame(idx, index=idx)
         tm.assert_frame_equal(result, expected)
 
-    @pytest.mark.parametrize("name", [None, "name"])
-    def test_index_map(self, name):
-        # see GH20990
-        count = 6
-        index = pd.date_range("2018-01-01", periods=count, freq="M", name=name).map(
-            lambda x: (x.year, x.month)
-        )
-        exp_index = pd.MultiIndex.from_product(
-            ((2018,), range(1, 7)), names=[name, name]
-        )
-        tm.assert_index_equal(index, exp_index)
-
     def test_split_non_utc(self):
         # GH 14042
         indices = pd.date_range("2016-01-01 00:00:00+0200", freq="S", periods=10)

pandas/tests/indexes/datetimes/test_indexing.py

@@ -518,6 +518,21 @@ def test_dti_contains_with_duplicates(self):
         ix = DatetimeIndex([d, d])
         assert d in ix
 
+    @pytest.mark.parametrize(
+        "vals",
+        [
+            [0, 1, 0],
+            [0, 0, -1],
+            [0, -1, -1],
+            ["2015", "2015", "2016"],
+            ["2015", "2015", "2014"],
+        ],
+    )
+    def test_contains_nonunique(self, vals):
+        # GH#9512
+        idx = DatetimeIndex(vals)
+        assert idx[0] in idx
+
 
 class TestGetIndexer:
     def test_get_indexer(self):

pandas/tests/indexes/datetimes/test_map.py

@@ -0,0 +1,41 @@
+import pytest
+
+from pandas import DatetimeIndex, Index, MultiIndex, Period, date_range
+import pandas._testing as tm
+
+
+class TestMap:
+    def test_map(self):
+        rng = date_range("1/1/2000", periods=10)
+
+        f = lambda x: x.strftime("%Y%m%d")
+        result = rng.map(f)
+        exp = Index([f(x) for x in rng], dtype="<U8")
+        tm.assert_index_equal(result, exp)
+
+    def test_map_fallthrough(self, capsys):
+        # GH#22067, check we don't get warnings about silently ignored errors
+        dti = date_range("2017-01-01", "2018-01-01", freq="B")
+
+        dti.map(lambda x: Period(year=x.year, month=x.month, freq="M"))
+
+        captured = capsys.readouterr()
+        assert captured.err == ""
+
+    def test_map_bug_1677(self):
+        index = DatetimeIndex(["2012-04-25 09:30:00.393000"])
+        f = index.asof
+
+        result = index.map(f)
+        expected = Index([f(index[0])])
+        tm.assert_index_equal(result, expected)
+
+    @pytest.mark.parametrize("name", [None, "name"])
+    def test_index_map(self, name):
+        # see GH#20990
+        count = 6
+        index = date_range("2018-01-01", periods=count, freq="M", name=name).map(
+            lambda x: (x.year, x.month)
+        )
+        exp_index = MultiIndex.from_product(((2018,), range(1, 7)), names=[name, name])
+        tm.assert_index_equal(index, exp_index)

pandas/tests/indexes/datetimes/test_ops.py

@@ -170,20 +170,6 @@ def test_value_counts_unique(self, tz_naive_fixture):
 
         tm.assert_index_equal(idx.unique(), exp_idx)
 
-    def test_nonunique_contains(self):
-        # GH 9512
-        for idx in map(
-            DatetimeIndex,
-            (
-                [0, 1, 0],
-                [0, 0, -1],
-                [0, -1, -1],
-                ["2015", "2015", "2016"],
-                ["2015", "2015", "2014"],
-            ),
-        ):
-            assert idx[0] in idx
-
     @pytest.mark.parametrize(
         "idx",
         [
@@ -432,10 +418,6 @@ def test_comparison(self):
         assert comp[11]
         assert not comp[9]
 
-    def test_pickle_unpickle(self):
-        unpickled = tm.round_trip_pickle(self.rng)
-        assert unpickled.freq is not None
-
     def test_copy(self):
         cp = self.rng.copy()
         repr(cp)
@@ -478,9 +460,5 @@ def test_copy(self):
         repr(cp)
         tm.assert_index_equal(cp, self.rng)
 
-    def test_pickle_unpickle(self):
-        unpickled = tm.round_trip_pickle(self.rng)
-        assert unpickled.freq is not None
-
     def test_equals(self):
         assert not self.rng.equals(list(self.rng))

pandas/tests/indexes/datetimes/test_pickle.py

@@ -0,0 +1,41 @@
+import pytest
+
+from pandas import NaT, date_range, to_datetime
+import pandas._testing as tm
+
+
+class TestPickle:
+    def test_pickle(self):
+        # GH#4606
+        idx = to_datetime(["2013-01-01", NaT, "2014-01-06"])
+        idx_p = tm.round_trip_pickle(idx)
+        assert idx_p[0] == idx[0]
+        assert idx_p[1] is NaT
+        assert idx_p[2] == idx[2]
+
+    def test_pickle_dont_infer_freq(self):
+        # GH##11002
+        # don't infer freq
+        idx = date_range("1750-1-1", "2050-1-1", freq="7D")
+        idx_p = tm.round_trip_pickle(idx)
+        tm.assert_index_equal(idx, idx_p)
+
+    def test_pickle_after_set_freq(self):
+        dti = date_range("20130101", periods=3, tz="US/Eastern", name="foo")
+        dti = dti._with_freq(None)
+
+        res = tm.round_trip_pickle(dti)
+        tm.assert_index_equal(res, dti)
+
+    def test_roundtrip_pickle_with_tz(self):
+        # GH#8367
+        # round-trip of timezone
+        index = date_range("20130101", periods=3, tz="US/Eastern", name="foo")
+        unpickled = tm.round_trip_pickle(index)
+        tm.assert_index_equal(index, unpickled)
+
+    @pytest.mark.parametrize("freq", ["B", "C"])
+    def test_pickle_unpickle(self, freq):
+        rng = date_range("2009-01-01", "2010-01-01", freq=freq)
+        unpickled = tm.round_trip_pickle(rng)
+        assert unpickled.freq == freq

pandas/tests/indexes/period/test_factorize.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+from pandas import PeriodIndex
+import pandas._testing as tm
+
+
+class TestFactorize:
+    def test_factorize(self):
+        idx1 = PeriodIndex(
+            ["2014-01", "2014-01", "2014-02", "2014-02", "2014-03", "2014-03"], freq="M"
+        )
+
+        exp_arr = np.array([0, 0, 1, 1, 2, 2], dtype=np.intp)
+        exp_idx = PeriodIndex(["2014-01", "2014-02", "2014-03"], freq="M")
+
+        arr, idx = idx1.factorize()
+        tm.assert_numpy_array_equal(arr, exp_arr)
+        tm.assert_index_equal(idx, exp_idx)
+
+        arr, idx = idx1.factorize(sort=True)
+        tm.assert_numpy_array_equal(arr, exp_arr)
+        tm.assert_index_equal(idx, exp_idx)
+
+        idx2 = PeriodIndex(
+            ["2014-03", "2014-03", "2014-02", "2014-01", "2014-03", "2014-01"], freq="M"
+        )
+
+        exp_arr = np.array([2, 2, 1, 0, 2, 0], dtype=np.intp)
+        arr, idx = idx2.factorize(sort=True)
+        tm.assert_numpy_array_equal(arr, exp_arr)
+        tm.assert_index_equal(idx, exp_idx)
+
+        exp_arr = np.array([0, 0, 1, 2, 0, 2], dtype=np.intp)
+        exp_idx = PeriodIndex(["2014-03", "2014-02", "2014-01"], freq="M")
+        arr, idx = idx2.factorize()
+        tm.assert_numpy_array_equal(arr, exp_arr)
+        tm.assert_index_equal(idx, exp_idx)

pandas/tests/indexes/period/test_indexing.py

@@ -764,6 +764,27 @@ def test_contains(self):
 
         assert p3 not in idx0
 
+    def test_contains_freq_mismatch(self):
+        rng = period_range("2007-01", freq="M", periods=10)
+
+        assert Period("2007-01", freq="M") in rng
+        assert not Period("2007-01", freq="D") in rng
+        assert not Period("2007-01", freq="2M") in rng
+
+    def test_contains_nat(self):
+        # see gh-13582
+        idx = period_range("2007-01", freq="M", periods=10)
+        assert NaT not in idx
+        assert None not in idx
+        assert float("nan") not in idx
+        assert np.nan not in idx
+
+        idx = PeriodIndex(["2011-01", "NaT", "2011-02"], freq="M")
+        assert NaT in idx
+        assert None in idx
+        assert float("nan") in idx
+        assert np.nan in idx
+
 
 class TestAsOfLocs:
     def test_asof_locs_mismatched_type(self):

pandas/tests/indexes/period/test_period.py

@@ -318,37 +318,6 @@ def test_period_reindex_with_object(
         expected = pd.Series(expected_values, index=object_index)
         tm.assert_series_equal(result, expected)
 
-    def test_factorize(self):
-        idx1 = PeriodIndex(
-            ["2014-01", "2014-01", "2014-02", "2014-02", "2014-03", "2014-03"], freq="M"
-        )
-
-        exp_arr = np.array([0, 0, 1, 1, 2, 2], dtype=np.intp)
-        exp_idx = PeriodIndex(["2014-01", "2014-02", "2014-03"], freq="M")
-
-        arr, idx = idx1.factorize()
-        tm.assert_numpy_array_equal(arr, exp_arr)
-        tm.assert_index_equal(idx, exp_idx)
-
-        arr, idx = idx1.factorize(sort=True)
-        tm.assert_numpy_array_equal(arr, exp_arr)
-        tm.assert_index_equal(idx, exp_idx)
-
-        idx2 = PeriodIndex(
-            ["2014-03", "2014-03", "2014-02", "2014-01", "2014-03", "2014-01"], freq="M"
-        )
-
-        exp_arr = np.array([2, 2, 1, 0, 2, 0], dtype=np.intp)
-        arr, idx = idx2.factorize(sort=True)
-        tm.assert_numpy_array_equal(arr, exp_arr)
-        tm.assert_index_equal(idx, exp_idx)
-
-        exp_arr = np.array([0, 0, 1, 2, 0, 2], dtype=np.intp)
-        exp_idx = PeriodIndex(["2014-03", "2014-02", "2014-01"], freq="M")
-        arr, idx = idx2.factorize()
-        tm.assert_numpy_array_equal(arr, exp_arr)
-        tm.assert_index_equal(idx, exp_idx)
-
     def test_is_(self):
         create_index = lambda: period_range(freq="A", start="1/1/2001", end="12/1/2009")
         index = create_index()
@@ -367,27 +336,6 @@ def test_is_(self):
         assert not index.is_(index - 2)
         assert not index.is_(index - 0)
 
-    def test_contains(self):
-        rng = period_range("2007-01", freq="M", periods=10)
-
-        assert Period("2007-01", freq="M") in rng
-        assert not Period("2007-01", freq="D") in rng
-        assert not Period("2007-01", freq="2M") in rng
-
-    def test_contains_nat(self):
-        # see gh-13582
-        idx = period_range("2007-01", freq="M", periods=10)
-        assert NaT not in idx
-        assert None not in idx
-        assert float("nan") not in idx
-        assert np.nan not in idx
-
-        idx = PeriodIndex(["2011-01", "NaT", "2011-02"], freq="M")
-        assert NaT in idx
-        assert None in idx
-        assert float("nan") in idx
-        assert np.nan in idx
-
     def test_periods_number_check(self):
         msg = (
             "Of the three parameters: start, end, and periods, exactly two "