Inline check(assert_type(...)) calls

Author: dpoznik

tests/test_frame.py

@@ -1960,25 +1960,41 @@ def test_types_rename_axis() -> None:
     df.index.name = "a"
     df.columns.name = "b"
 
-    for _df in [
-        # Rename axes with `mapper` and `axis`
-        df.rename_axis("A"),
-        df.rename_axis(["A"]),
-        df.rename_axis(None),
-        df.rename_axis("B", axis=1),
-        df.rename_axis(["B"], axis=1),
-        df.rename_axis(None, axis=1),
-        # Rename axes with `index` and `columns`
-        df.rename_axis(index="A", columns="B"),
-        df.rename_axis(index=["A"], columns=["B"]),
-        df.rename_axis(index={"a": "A"}, columns={"b": "B"}),
-        df.rename_axis(
-            index=lambda name: name.upper(),
-            columns=lambda name: name.upper(),
+    # Rename axes with `mapper` and `axis`
+    check(assert_type(df.rename_axis("A"), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.rename_axis(["A"]), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.rename_axis(None), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.rename_axis("B", axis=1), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.rename_axis(["B"], axis=1), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.rename_axis(None, axis=1), pd.DataFrame), pd.DataFrame)
+
+    # Rename axes with `index` and `columns`
+    check(
+        assert_type(df.rename_axis(index="A", columns="B"), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(df.rename_axis(index=["A"], columns=["B"]), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(df.rename_axis(index={"a": "A"}, columns={"b": "B"}), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(
+            df.rename_axis(
+                index=lambda name: name.upper(),
+                columns=lambda name: name.upper(),
+            ),
+            pd.DataFrame,
         ),
-        df.rename_axis(index=None, columns=None),
-    ]:
-        check(assert_type(_df, pd.DataFrame), pd.DataFrame)
+        pd.DataFrame,
+    )
+    check(
+        assert_type(df.rename_axis(index=None, columns=None), pd.DataFrame),
+        pd.DataFrame,
+    )
 
 
 def test_types_eq() -> None:

tests/test_series.py

@@ -1124,19 +1124,20 @@ def test_types_rename_axis() -> None:
     s = pd.Series([1, 2, 3])
     s.index.name = "a"
 
-    for _s in [
-        # Rename index with `mapper`
-        s.rename_axis("A"),
-        s.rename_axis(["A"]),
-        s.rename_axis(None),
-        # Rename index with `index`
-        s.rename_axis(index="A"),
-        s.rename_axis(index=["A"]),
-        s.rename_axis(index={"a": "A"}),
-        s.rename_axis(index=lambda name: name.upper()),
-        s.rename_axis(index=None),
-    ]:
-        check(assert_type(_s, pd.Series[int]), pd.Series[int])
+    # Rename index with `mapper`
+    check(assert_type(s.rename_axis("A"), "pd.Series[int]"), pd.Series)
+    check(assert_type(s.rename_axis(["A"]), "pd.Series[int]"), pd.Series)
+    check(assert_type(s.rename_axis(None), "pd.Series[int]"), pd.Series)
+
+    # Rename index with `index`
+    check(assert_type(s.rename_axis(index="A"), "pd.Series[int]"), pd.Series)
+    check(assert_type(s.rename_axis(index=["A"]), "pd.Series[int]"), pd.Series)
+    check(assert_type(s.rename_axis(index={"a": "A"}), "pd.Series[int]"), pd.Series)
+    check(
+        assert_type(s.rename_axis(index=lambda name: name.upper()), "pd.Series[int]"),
+        pd.Series,
+    )
+    check(assert_type(s.rename_axis(index=None), "pd.Series[int]"), pd.Series)
 
 
 def test_types_values() -> None: