Add flake8 extension: comprehensions (#4200)

py-polars/build.requirements.txt

@@ -21,6 +21,7 @@ mypy==0.961
 ghp-import==2.1.0
 flake8==4.0.1
 flake8-bugbear==22.7.1
+flake8-comprehensions==3.10.0
 flake8-docstrings==1.6.0
 sphinx==4.2.0
 pydata-sphinx-theme==0.6.3

py-polars/polars/_html.py

@@ -159,7 +159,7 @@ def write_style(self) -> None:
         if not in_vscode_notebook():
             element_props.append(("td", "line-height", "95%"))
 
-        template_mid = "\n\n".join(map(lambda t: template_select % t, element_props))
+        template_mid = "\n\n".join(template_select % t for t in element_props)
         template = dedent("\n".join((template_first, template_mid, template_last)))
         self.write(template)
 

py-polars/polars/internals/frame.py

@@ -4499,7 +4499,7 @@ def get_columns(self) -> list[pli.Series]:
         ]]
 
         """
-        return list(map(lambda s: pli.wrap_s(s), self._df.get_columns()))
+        return [pli.wrap_s(s) for s in self._df.get_columns()]
 
     def get_column(self, name: str) -> pli.Series:
         """
@@ -4941,7 +4941,7 @@ def partition_by(
             groups = [groups]
 
         if as_dict:
-            out: dict[Any, DF] = dict()
+            out: dict[Any, DF] = {}
             if len(groups) == 1:
                 for _df in self._df.partition_by(groups, maintain_order):
                     df = self._from_pydf(_df)

py-polars/tests/io/test_csv.py

@@ -585,7 +585,7 @@ def test_different_eol_char() -> None:
 
 
 def test_csv_write_escape_newlines() -> None:
-    df = pl.DataFrame(dict(escape=["n\nn"]))
+    df = pl.DataFrame({"escape": ["n\nn"]})
     f = io.BytesIO()
     df.write_csv(f)
     f.seek(0)

py-polars/tests/test_apply.py

@@ -125,9 +125,11 @@ def test_apply_struct() -> None:
 
 def test_apply_numpy_out_3057() -> None:
     df = pl.DataFrame(
-        dict(
-            id=[0, 0, 0, 1, 1, 1], t=[2.0, 4.3, 5, 10, 11, 14], y=[0.0, 1, 1.3, 2, 3, 4]
-        )
+        {
+            "id": [0, 0, 0, 1, 1, 1],
+            "t": [2.0, 4.3, 5, 10, 11, 14],
+            "y": [0.0, 1, 1.3, 2, 3, 4],
+        }
     )
 
     assert (

py-polars/tests/test_df.py

@@ -1850,13 +1850,13 @@ class SubClassedDataFrame(pl.DataFrame):
 
 def test_explode_empty() -> None:
     df = (
-        pl.DataFrame(dict(x=["a", "a", "b", "b"], y=[1, 1, 2, 2]))
+        pl.DataFrame({"x": ["a", "a", "b", "b"], "y": [1, 1, 2, 2]})
         .groupby("x")
         .agg(pl.col("y").take([]))
     )
     assert df.explode("y").shape == (0, 2)
 
-    df = pl.DataFrame(dict(x=["1", "2", "4"], y=[["a", "b", "c"], ["d"], []]))
+    df = pl.DataFrame({"x": ["1", "2", "4"], "y": [["a", "b", "c"], ["d"], []]})
     assert df.explode("y").frame_equal(
         pl.DataFrame({"x": ["1", "1", "1", "2", "4"], "y": ["a", "b", "c", "d", None]})
     )

py-polars/tests/test_errors.py

@@ -13,12 +13,12 @@ def test_error_on_empty_groupby() -> None:
     with pytest.raises(
         pl.ComputeError, match="expected keys in groupby operation, got nothing"
     ):
-        pl.DataFrame(dict(x=[0, 0, 1, 1])).groupby([]).agg(pl.count())
+        pl.DataFrame({"x": [0, 0, 1, 1]}).groupby([]).agg(pl.count())
 
 
 def test_error_on_reducing_map() -> None:
     df = pl.DataFrame(
-        dict(id=[0, 0, 0, 1, 1, 1], t=[2, 4, 5, 10, 11, 14], y=[0, 1, 1, 2, 3, 4])
+        {"id": [0, 0, 0, 1, 1, 1], "t": [2, 4, 5, 10, 11, 14], "y": [0, 1, 1, 2, 3, 4]}
     )
 
     with pytest.raises(

py-polars/tests/test_exprs.py

@@ -129,7 +129,7 @@ def test_wildcard_expansion() -> None:
 
 
 def test_split_exact() -> None:
-    df = pl.DataFrame(dict(x=["a_a", None, "b", "c_c"]))
+    df = pl.DataFrame({"x": ["a_a", None, "b", "c_c"]})
     out = df.select([pl.col("x").str.split_exact("_", 2, inclusive=False)]).unnest("x")
 
     expected = pl.DataFrame(

py-polars/tests/test_window.py

@@ -97,13 +97,13 @@ def test_window_function_cache() -> None:
 
 
 def test_arange_no_rows() -> None:
-    df = pl.DataFrame(dict(x=[5, 5, 4, 4, 2, 2]))
+    df = pl.DataFrame({"x": [5, 5, 4, 4, 2, 2]})
     out = df.with_column(pl.arange(0, pl.count()).over("x"))  # type: ignore[union-attr]
     assert out.frame_equal(
         pl.DataFrame({"x": [5, 5, 4, 4, 2, 2], "literal": [0, 1, 0, 1, 0, 1]})
     )
 
-    df = pl.DataFrame(dict(x=[]))
+    df = pl.DataFrame({"x": []})
     out = df.with_column(pl.arange(0, pl.count()).over("x"))  # type: ignore[union-attr]
     assert out.frame_equal(pl.DataFrame({"x": [], "literal": []}))