feat(python): improved support for use of file-like objects with `Dat…

…aFrame` "write" methods (#12113)

py-polars/polars/dataframe/frame.py

@@ -2358,8 +2358,8 @@ def write_json(
         Parameters
         ----------
         file
-            File path to which the result should be written. If set to `None`
-            (default), the output is returned as a string instead.
+            File path or writeable file-like object to which the result will be written.
+            If set to `None` (default), the output is returned as a string instead.
         pretty
             Pretty serialize json.
         row_oriented
@@ -2415,8 +2415,8 @@ def write_ndjson(self, file: IOBase | str | Path | None = None) -> str | None:
         Parameters
         ----------
         file
-            File path to which the result should be written. If set to `None`
-            (default), the output is returned as a string instead.
+            File path or writeable file-like object to which the result will be written.
+            If set to `None` (default), the output is returned as a string instead.
 
         Examples
         --------
@@ -2508,8 +2508,8 @@ def write_csv(
         Parameters
         ----------
         file
-            File path to which the result should be written. If set to `None`
-            (default), the output is returned as a string instead.
+            File path or writeable file-like object to which the result will be written.
+            If set to `None` (default), the output is returned as a string instead.
         has_header
             Whether to include header in the CSV output.
         separator
@@ -2616,7 +2616,7 @@ def write_avro(
         Parameters
         ----------
         file
-            File path to which the file should be written.
+            File path or writeable file-like object to which the data will be written.
         compression : {'uncompressed', 'snappy', 'deflate'}
             Compression method. Defaults to "uncompressed".
 
@@ -3178,8 +3178,8 @@ def write_ipc(
         Parameters
         ----------
         file
-            Path to which the IPC data should be written. If set to
-            `None`, the output is returned as a BytesIO object.
+            Path or writeable file-like object to which the IPC data will be
+            written. If set to `None`, the output is returned as a BytesIO object.
         compression : {'uncompressed', 'lz4', 'zstd'}
             Compression method. Defaults to "uncompressed".
 
@@ -3239,8 +3239,8 @@ def write_ipc_stream(
         Parameters
         ----------
         file
-            Path to which the IPC record batch data should be written. If set to
-            `None`, the output is returned as a BytesIO object.
+            Path or writeable file-like object to which the IPC record batch data will
+            be written. If set to `None`, the output is returned as a BytesIO object.
         compression : {'uncompressed', 'lz4', 'zstd'}
             Compression method. Defaults to "uncompressed".
 
@@ -3288,7 +3288,7 @@ def write_parquet(
         Parameters
         ----------
         file
-            File path to which the file should be written.
+            File path or writeable file-like object to which the result will be written.
         compression : {'lz4', 'uncompressed', 'snappy', 'gzip', 'lzo', 'brotli', 'zstd'}
             Choose "zstd" for good compression performance.
             Choose "lz4" for fast compression/decompression.

py-polars/src/file.rs

@@ -222,7 +222,7 @@ pub fn get_either_file(py_f: PyObject, truncate: bool) -> PyResult<EitherRustPyt
             };
             Ok(EitherRustPythonFile::Rust(f))
         } else {
-            let f = PyFileLikeObject::with_requirements(py_f, true, true, true)?;
+            let f = PyFileLikeObject::with_requirements(py_f, !truncate, truncate, !truncate)?;
             Ok(EitherRustPythonFile::Py(f))
         }
     })

py-polars/tests/unit/io/test_csv.py

@@ -1603,3 +1603,22 @@ def test_provide_schema() -> None:
         "B": [None, "ragged", None],
         "C": [None, None, None],
     }
+
+
+def test_custom_writeable_object() -> None:
+    df = pl.DataFrame({"a": [10, 20, 30], "b": ["x", "y", "z"]})
+
+    class CustomBuffer:
+        writes: list[bytes]
+
+        def __init__(self) -> None:
+            self.writes = []
+
+        def write(self, data: bytes) -> int:
+            self.writes.append(data)
+            return len(data)
+
+    buf = CustomBuffer()
+    df.write_csv(buf)  # type: ignore[call-overload]
+
+    assert b"".join(buf.writes) == b"a,b\n10,x\n20,y\n30,z\n"