use IO[bytes] instead of BytesIO

pola-rs · Sep 10, 2024 · d74341e · d74341e
1 parent 5ccb238
commit d74341e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py
@@ -3638,7 +3638,7 @@ def write_ipc_stream(
 
     def write_parquet(
         self,
-        file: str | Path | BytesIO,
+        file: str | Path | IO[bytes],
         *,
         compression: ParquetCompression = "zstd",
         compression_level: int | None = None,

diff --git a/py-polars/tests/unit/io/test_parquet.py b/py-polars/tests/unit/io/test_parquet.py
@@ -1891,3 +1891,15 @@ def test_concat_multiple_inmem() -> None:
 
     assert_frame_equal(pl.read_parquet([fb, gb]), dfs)
     assert_frame_equal(pl.read_parquet([fb, gb], use_pyarrow=True), dfs)
+
+@pytest.mark.write_disk
+def test_write_binary_open_file(tmp_path):
+    df = pl.DataFrame({"a": [1, 2, 3]})
+
+    path = tmp_path / "test.parquet"
+
+    with path.open("wb") as f_write:
+        df.write_parquet(f_write)
+
+    out = pl.read_parquet(path)
+    assert_frame_equal(out, df)