Shift test functions up

Author: ryankarlos

pandas/tests/groupby/test_transform.py

@@ -356,6 +356,84 @@ def test_transform_transformation_func(transformation_func):
     tm.assert_frame_equal(result, expected)
 
 
+def test_groupby_corrwith(transformation_func, df_for_transformation_func):
+
+    # GH 27905
+    df = df_for_transformation_func.copy()
+    g = df.groupby("A")
+
+    if transformation_func == "corrwith":
+        op = lambda x: getattr(x, transformation_func)(df)
+        result = op(g)
+        expected = pd.DataFrame(dict(B=[1, np.nan, np.nan], A=[np.nan] * 3))
+        expected.index = pd.Index([121, 231, 676], name="A")
+        tm.assert_frame_equal(result, expected)
+
+
+def test_groupby_transform_nan(transformation_func, df_for_transformation_func):
+
+    # GH 27905
+    df = df_for_transformation_func.copy()
+    g = df.groupby("A")
+
+    if transformation_func == "fillna":
+
+        df["B"] = [1, np.nan, np.nan, 3, np.nan, 3, 4]
+        result = g.transform(transformation_func, value=1)
+        expected = pd.DataFrame({"B": [1.0, 1.0, 1.0, 3.0, 1.0, 3.0, 4.0]})
+        tm.assert_frame_equal(result, expected)
+        op = lambda x: getattr(x, transformation_func)(1)
+        result = op(g)
+        tm.assert_frame_equal(result, expected)
+
+
+def test_groupby_tshift(transformation_func, df_for_transformation_func):
+
+    # GH 27905
+    df = df_for_transformation_func.copy()
+    dt_periods = pd.date_range("2013-11-03", periods=7, freq="D")
+    df["C"] = dt_periods
+    g = df.set_index("C").groupby("A")
+
+    if transformation_func == "tshift":
+
+        op = lambda x: getattr(x, transformation_func)(2, "D")
+        result = op(g)
+        df["C"] = dt_periods + dt_periods.freq * 2
+        expected = df
+        tm.assert_frame_equal(
+            result.reset_index().reindex(columns=["A", "B", "C"]), expected
+        )
+
+
+def test_check_original_and_transformed_index(transformation_func):
+
+    # GH 27905
+    df = DataFrame(
+        {
+            "A": [121, 121, 121, 121, 231, 231, 676],
+            "B": [1.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0],
+        }
+    )
+
+    df = DataFrame({"A": [0, 0, 0, 1, 1, 1], "B": [0, 1, 2, 3, 4, 5]})
+    g = df.groupby("A")
+
+    if transformation_func in [
+        "cummax",
+        "cummin",
+        "cumprod",
+        "cumsum",
+        "diff",
+        "ffill",
+        "pct_change",
+        "rank",
+        "shift",
+    ]:
+        result = g.transform(transformation_func)
+        tm.assert_index_equal(result.index, df.index)
+
+
 def test_transform_select_columns(df):
     f = lambda x: x.mean()
     result = df.groupby("A")[["C", "D"]].transform(f)
@@ -1205,82 +1283,4 @@ def test_transform_lambda_indexing():
             names=["A", "B"],
         ),
     )
-    tm.assert_frame_equal(result, expected)
-
-
-def test_groupby_corrwith(transformation_func, df_for_transformation_func):
-
-    # GH 27905
-    df = df_for_transformation_func.copy()
-    g = df.groupby("A")
-
-    if transformation_func == "corrwith":
-        op = lambda x: getattr(x, transformation_func)(df)
-        result = op(g)
-        expected = pd.DataFrame(dict(B=[1, np.nan, np.nan], A=[np.nan] * 3))
-        expected.index = pd.Index([121, 231, 676], name="A")
-        tm.assert_frame_equal(result, expected)
-
-
-def test_groupby_transform_nan(transformation_func, df_for_transformation_func):
-
-    # GH 27905
-    df = df_for_transformation_func.copy()
-    g = df.groupby("A")
-
-    if transformation_func == "fillna":
-
-        df["B"] = [1, np.nan, np.nan, 3, np.nan, 3, 4]
-        result = g.transform(transformation_func, value=1)
-        expected = pd.DataFrame({"B": [1.0, 1.0, 1.0, 3.0, 1.0, 3.0, 4.0]})
-        tm.assert_frame_equal(result, expected)
-        op = lambda x: getattr(x, transformation_func)(1)
-        result = op(g)
-        tm.assert_frame_equal(result, expected)
-
-
-def test_groupby_tshift(transformation_func, df_for_transformation_func):
-
-    # GH 27905
-    df = df_for_transformation_func.copy()
-    dt_periods = pd.date_range("2013-11-03", periods=7, freq="D")
-    df["C"] = dt_periods
-    g = df.set_index("C").groupby("A")
-
-    if transformation_func == "tshift":
-
-        op = lambda x: getattr(x, transformation_func)(2, "D")
-        result = op(g)
-        df["C"] = dt_periods + dt_periods.freq * 2
-        expected = df
-        tm.assert_frame_equal(
-            result.reset_index().reindex(columns=["A", "B", "C"]), expected
-        )
-
-
-def test_check_original_and_transformed_index(transformation_func):
-
-    # GH 27905
-    df = DataFrame(
-        {
-            "A": [121, 121, 121, 121, 231, 231, 676],
-            "B": [1.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0],
-        }
-    )
-
-    df = DataFrame({"A": [0, 0, 0, 1, 1, 1], "B": [0, 1, 2, 3, 4, 5]})
-    g = df.groupby("A")
-
-    if transformation_func in [
-        "cummax",
-        "cummin",
-        "cumprod",
-        "cumsum",
-        "diff",
-        "ffill",
-        "pct_change",
-        "rank",
-        "shift",
-    ]:
-        result = g.transform(transformation_func)
-        tm.assert_index_equal(result.index, df.index)
+    tm.assert_frame_equal(result, expected)