TST/REF: share tests across Series/DataFrame (#37616)

Author: jbrockmendel

pandas/tests/frame/methods/test_asof.py

@@ -96,12 +96,16 @@ def test_missing(self, date_range_frame):
         result = df.asof("1989-12-31")
         assert isinstance(result.name, Period)
 
+    def test_asof_all_nans(self, frame_or_series):
+        # GH 15713
+        # DataFrame/Series is all nans
+        result = frame_or_series([np.nan]).asof([0])
+        expected = frame_or_series([np.nan])
+        tm.assert_equal(result, expected)
+
     def test_all_nans(self, date_range_frame):
         # GH 15713
         # DataFrame is all nans
-        result = DataFrame([np.nan]).asof([0])
-        expected = DataFrame([np.nan])
-        tm.assert_frame_equal(result, expected)
 
         # testing non-default indexes, multiple inputs
         N = 150

pandas/tests/frame/methods/test_droplevel.py

@@ -1,23 +1,32 @@
+import pytest
+
 from pandas import DataFrame, Index, MultiIndex
 import pandas._testing as tm
 
 
 class TestDropLevel:
-    def test_droplevel(self):
+    def test_droplevel(self, frame_or_series):
         # GH#20342
-        df = DataFrame([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
-        df = df.set_index([0, 1]).rename_axis(["a", "b"])
-        df.columns = MultiIndex.from_tuples(
+        cols = MultiIndex.from_tuples(
             [("c", "e"), ("d", "f")], names=["level_1", "level_2"]
         )
+        mi = MultiIndex.from_tuples([(1, 2), (5, 6), (9, 10)], names=["a", "b"])
+        df = DataFrame([[3, 4], [7, 8], [11, 12]], index=mi, columns=cols)
+        if frame_or_series is not DataFrame:
+            df = df.iloc[:, 0]
 
         # test that dropping of a level in index works
         expected = df.reset_index("a", drop=True)
         result = df.droplevel("a", axis="index")
-        tm.assert_frame_equal(result, expected)
+        tm.assert_equal(result, expected)
 
-        # test that dropping of a level in columns works
-        expected = df.copy()
-        expected.columns = Index(["c", "d"], name="level_1")
-        result = df.droplevel("level_2", axis="columns")
-        tm.assert_frame_equal(result, expected)
+        if frame_or_series is DataFrame:
+            # test that dropping of a level in columns works
+            expected = df.copy()
+            expected.columns = Index(["c", "d"], name="level_1")
+            result = df.droplevel("level_2", axis="columns")
+            tm.assert_equal(result, expected)
+        else:
+            # test that droplevel raises ValueError on axis != 0
+            with pytest.raises(ValueError, match="No axis named columns"):
+                df.droplevel(1, axis="columns")

pandas/tests/frame/methods/test_first_and_last.py

@@ -8,56 +8,64 @@
 
 
 class TestFirst:
-    def test_first_subset(self):
+    def test_first_subset(self, frame_or_series):
         ts = tm.makeTimeDataFrame(freq="12h")
+        if frame_or_series is not DataFrame:
+            ts = ts["A"]
         result = ts.first("10d")
         assert len(result) == 20
 
         ts = tm.makeTimeDataFrame(freq="D")
+        if frame_or_series is not DataFrame:
+            ts = ts["A"]
         result = ts.first("10d")
         assert len(result) == 10
 
         result = ts.first("3M")
         expected = ts[:"3/31/2000"]
-        tm.assert_frame_equal(result, expected)
+        tm.assert_equal(result, expected)
 
         result = ts.first("21D")
         expected = ts[:21]
-        tm.assert_frame_equal(result, expected)
+        tm.assert_equal(result, expected)
 
         result = ts[:0].first("3M")
-        tm.assert_frame_equal(result, ts[:0])
+        tm.assert_equal(result, ts[:0])
 
-    def test_first_raises(self):
+    def test_first_last_raises(self, frame_or_series):
         # GH#20725
-        df = DataFrame([[1, 2, 3], [4, 5, 6]])
+        obj = DataFrame([[1, 2, 3], [4, 5, 6]])
+        if frame_or_series is not DataFrame:
+            obj = obj[0]
+
         msg = "'first' only supports a DatetimeIndex index"
         with pytest.raises(TypeError, match=msg):  # index is not a DatetimeIndex
-            df.first("1D")
+            obj.first("1D")
+
+        msg = "'last' only supports a DatetimeIndex index"
+        with pytest.raises(TypeError, match=msg):  # index is not a DatetimeIndex
+            obj.last("1D")
 
-    def test_last_subset(self):
+    def test_last_subset(self, frame_or_series):
         ts = tm.makeTimeDataFrame(freq="12h")
+        if frame_or_series is not DataFrame:
+            ts = ts["A"]
         result = ts.last("10d")
         assert len(result) == 20
 
         ts = tm.makeTimeDataFrame(nper=30, freq="D")
+        if frame_or_series is not DataFrame:
+            ts = ts["A"]
         result = ts.last("10d")
         assert len(result) == 10
 
         result = ts.last("21D")
         expected = ts["2000-01-10":]
-        tm.assert_frame_equal(result, expected)
+        tm.assert_equal(result, expected)
 
         result = ts.last("21D")
         expected = ts[-21:]
-        tm.assert_frame_equal(result, expected)
+        tm.assert_equal(result, expected)
 
         result = ts[:0].last("3M")
-        tm.assert_frame_equal(result, ts[:0])
-
-    def test_last_raises(self):
-        # GH20725
-        df = DataFrame([[1, 2, 3], [4, 5, 6]])
-        msg = "'last' only supports a DatetimeIndex index"
-        with pytest.raises(TypeError, match=msg):  # index is not a DatetimeIndex
-            df.last("1D")
+        tm.assert_equal(result, ts[:0])

pandas/tests/frame/methods/test_head_tail.py

@@ -48,6 +48,9 @@ def test_head_tail(float_frame):
     tm.assert_frame_equal(df.tail(0), df[0:0])
     tm.assert_frame_equal(df.head(-1), df.iloc[:-1])
     tm.assert_frame_equal(df.tail(-1), df.iloc[1:])
+
+
+def test_head_tail_empty():
     # test empty dataframe
     empty_df = DataFrame()
     tm.assert_frame_equal(empty_df.tail(), empty_df)

pandas/tests/frame/methods/test_truncate.py

@@ -2,12 +2,15 @@
 import pytest
 
 import pandas as pd
+from pandas import DataFrame, Series, date_range
 import pandas._testing as tm
 
 
 class TestDataFrameTruncate:
-    def test_truncate(self, datetime_frame):
+    def test_truncate(self, datetime_frame, frame_or_series):
         ts = datetime_frame[::3]
+        if frame_or_series is Series:
+            ts = ts.iloc[:, 0]
 
         start, end = datetime_frame.index[3], datetime_frame.index[6]
 
@@ -16,34 +19,41 @@ def test_truncate(self, datetime_frame):
 
         # neither specified
         truncated = ts.truncate()
-        tm.assert_frame_equal(truncated, ts)
+        tm.assert_equal(truncated, ts)
 
         # both specified
         expected = ts[1:3]
 
         truncated = ts.truncate(start, end)
-        tm.assert_frame_equal(truncated, expected)
+        tm.assert_equal(truncated, expected)
 
         truncated = ts.truncate(start_missing, end_missing)
-        tm.assert_frame_equal(truncated, expected)
+        tm.assert_equal(truncated, expected)
 
         # start specified
         expected = ts[1:]
 
         truncated = ts.truncate(before=start)
-        tm.assert_frame_equal(truncated, expected)
+        tm.assert_equal(truncated, expected)
 
         truncated = ts.truncate(before=start_missing)
-        tm.assert_frame_equal(truncated, expected)
+        tm.assert_equal(truncated, expected)
 
         # end specified
         expected = ts[:3]
 
         truncated = ts.truncate(after=end)
-        tm.assert_frame_equal(truncated, expected)
+        tm.assert_equal(truncated, expected)
 
         truncated = ts.truncate(after=end_missing)
-        tm.assert_frame_equal(truncated, expected)
+        tm.assert_equal(truncated, expected)
+
+        # corner case, empty series/frame returned
+        truncated = ts.truncate(after=ts.index[0] - ts.index.freq)
+        assert len(truncated) == 0
+
+        truncated = ts.truncate(before=ts.index[-1] + ts.index.freq)
+        assert len(truncated) == 0
 
         msg = "Truncate: 2000-01-06 00:00:00 must be after 2000-02-04 00:00:00"
         with pytest.raises(ValueError, match=msg):
@@ -57,25 +67,35 @@ def test_truncate_copy(self, datetime_frame):
         truncated.values[:] = 5.0
         assert not (datetime_frame.values[5:11] == 5).any()
 
-    def test_truncate_nonsortedindex(self):
+    def test_truncate_nonsortedindex(self, frame_or_series):
         # GH#17935
 
-        df = pd.DataFrame({"A": ["a", "b", "c", "d", "e"]}, index=[5, 3, 2, 9, 0])
+        obj = DataFrame({"A": ["a", "b", "c", "d", "e"]}, index=[5, 3, 2, 9, 0])
+        if frame_or_series is Series:
+            obj = obj["A"]
+
         msg = "truncate requires a sorted index"
         with pytest.raises(ValueError, match=msg):
-            df.truncate(before=3, after=9)
+            obj.truncate(before=3, after=9)
+
+    def test_sort_values_nonsortedindex(self):
+        # TODO: belongs elsewhere?
 
-        rng = pd.date_range("2011-01-01", "2012-01-01", freq="W")
-        ts = pd.DataFrame(
+        rng = date_range("2011-01-01", "2012-01-01", freq="W")
+        ts = DataFrame(
             {"A": np.random.randn(len(rng)), "B": np.random.randn(len(rng))}, index=rng
         )
+
         msg = "truncate requires a sorted index"
         with pytest.raises(ValueError, match=msg):
             ts.sort_values("A", ascending=False).truncate(
                 before="2011-11", after="2011-12"
             )
 
-        df = pd.DataFrame(
+    def test_truncate_nonsortedindex_axis1(self):
+        # GH#17935
+
+        df = DataFrame(
             {
                 3: np.random.randn(5),
                 20: np.random.randn(5),
@@ -93,27 +113,34 @@ def test_truncate_nonsortedindex(self):
         [(1, 2, [2, 1]), (None, 2, [2, 1, 0]), (1, None, [3, 2, 1])],
     )
     @pytest.mark.parametrize("klass", [pd.Int64Index, pd.DatetimeIndex])
-    def test_truncate_decreasing_index(self, before, after, indices, klass):
+    def test_truncate_decreasing_index(
+        self, before, after, indices, klass, frame_or_series
+    ):
         # https://github.com/pandas-dev/pandas/issues/33756
         idx = klass([3, 2, 1, 0])
         if klass is pd.DatetimeIndex:
             before = pd.Timestamp(before) if before is not None else None
             after = pd.Timestamp(after) if after is not None else None
             indices = [pd.Timestamp(i) for i in indices]
-        values = pd.DataFrame(range(len(idx)), index=idx)
+        values = frame_or_series(range(len(idx)), index=idx)
         result = values.truncate(before=before, after=after)
         expected = values.loc[indices]
-        tm.assert_frame_equal(result, expected)
+        tm.assert_equal(result, expected)
 
-    def test_truncate_multiindex(self):
+    def test_truncate_multiindex(self, frame_or_series):
         # GH 34564
         mi = pd.MultiIndex.from_product([[1, 2, 3, 4], ["A", "B"]], names=["L1", "L2"])
-        s1 = pd.DataFrame(range(mi.shape[0]), index=mi, columns=["col"])
+        s1 = DataFrame(range(mi.shape[0]), index=mi, columns=["col"])
+        if frame_or_series is Series:
+            s1 = s1["col"]
+
         result = s1.truncate(before=2, after=3)
 
-        df = pd.DataFrame.from_dict(
+        df = DataFrame.from_dict(
             {"L1": [2, 2, 3, 3], "L2": ["A", "B", "A", "B"], "col": [2, 3, 4, 5]}
         )
         expected = df.set_index(["L1", "L2"])
+        if frame_or_series is Series:
+            expected = expected["col"]
 
-        tm.assert_frame_equal(result, expected)
+        tm.assert_equal(result, expected)

pandas/tests/frame/methods/test_tz_convert.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 
-from pandas import DataFrame, Index, MultiIndex, Series, date_range
+from pandas import DataFrame, Index, MultiIndex, date_range
 import pandas._testing as tm
 
 
@@ -89,17 +89,16 @@ def test_tz_convert_and_localize(self, fn):
             df = DataFrame(index=l0)
             df = getattr(df, fn)("US/Pacific", level=1)
 
-    @pytest.mark.parametrize("klass", [Series, DataFrame])
     @pytest.mark.parametrize("copy", [True, False])
-    def test_tz_convert_copy_inplace_mutate(self, copy, klass):
+    def test_tz_convert_copy_inplace_mutate(self, copy, frame_or_series):
         # GH#6326
-        obj = klass(
+        obj = frame_or_series(
             np.arange(0, 5),
             index=date_range("20131027", periods=5, freq="1H", tz="Europe/Berlin"),
         )
         orig = obj.copy()
         result = obj.tz_convert("UTC", copy=copy)
-        expected = klass(np.arange(0, 5), index=obj.index.tz_convert("UTC"))
+        expected = frame_or_series(np.arange(0, 5), index=obj.index.tz_convert("UTC"))
         tm.assert_equal(result, expected)
         tm.assert_equal(obj, orig)
         assert result.index is not obj.index

pandas/tests/frame/methods/test_tz_localize.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 
-from pandas import DataFrame, Series, date_range
+from pandas import DataFrame, date_range
 import pandas._testing as tm
 
 
@@ -23,16 +23,15 @@ def test_frame_tz_localize(self):
         assert result.columns.tz.zone == "UTC"
         tm.assert_frame_equal(result, expected.T)
 
-    @pytest.mark.parametrize("klass", [Series, DataFrame])
     @pytest.mark.parametrize("copy", [True, False])
-    def test_tz_localize_copy_inplace_mutate(self, copy, klass):
+    def test_tz_localize_copy_inplace_mutate(self, copy, frame_or_series):
         # GH#6326
-        obj = klass(
+        obj = frame_or_series(
             np.arange(0, 5), index=date_range("20131027", periods=5, freq="1H", tz=None)
         )
         orig = obj.copy()
         result = obj.tz_localize("UTC", copy=copy)
-        expected = klass(
+        expected = frame_or_series(
             np.arange(0, 5),
             index=date_range("20131027", periods=5, freq="1H", tz="UTC"),
         )

pandas/tests/series/methods/test_asof.py

@@ -161,9 +161,6 @@ def test_errors(self):
     def test_all_nans(self):
         # GH 15713
         # series is all nans
-        result = Series([np.nan]).asof([0])
-        expected = Series([np.nan])
-        tm.assert_series_equal(result, expected)
 
         # testing non-default indexes
         N = 50