add axis test cases & refactor

Author: Ritsuki Yamada

tests/test_frame.py

@@ -2984,19 +2984,15 @@ def test_iloc_tuple() -> None:
     df = df.iloc[0:2,]
 
 
-def test_take_list() -> None:
-    df = pd.DataFrame({"a": [1, 2, 3]})
+def test_take() -> None:
+    df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
     check(assert_type(df.take([0, 1]), pd.DataFrame), pd.DataFrame)
-
-
-def test_take_list_npint() -> None:
-    df = pd.DataFrame({"a": [1, 2, 3]})
     check(assert_type(df.take([np.int64(0), np.int64(1)]), pd.DataFrame), pd.DataFrame)
-
-
-def test_take_ndarray() -> None:
-    df = pd.DataFrame({"a": [1, 2, 3]})
     check(assert_type(df.take(np.array([0, 1])), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.take([0, 1], "index"), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.take([0, 1], 0), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.take([0, 1], "columns"), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.take([0, 1], 1), pd.DataFrame), pd.DataFrame)
 
 
 def test_set_columns() -> None:

tests/test_series.py

@@ -1715,21 +1715,13 @@ def test_iloc_getitem_ndarray() -> None:
     check(assert_type(values_s.iloc[indices_u64], pd.Series), pd.Series)
 
 
-def test_take_list() -> None:
+def test_take() -> None:
     s = pd.Series(np.arange(10), name="a")
     check(assert_type(s.take([0, 1]), pd.Series), pd.Series)
-
-
-def test_take_list_npint() -> None:
-    s = pd.Series(np.arange(10), name="a")
     check(
         assert_type(s.take([np.int64(0), np.int64(1)]), pd.Series),
         pd.Series,
     )
-
-
-def test_take_ndarray() -> None:
-    s = pd.Series(np.arange(10), name="a")
     check(assert_type(s.take(np.array([0, 1])), pd.Series), pd.Series)