csv writer quote if string contains new line char (#4134)

pola-rs · Jul 24, 2022 · a006a14 · a006a14
1 parent 8d21c5b
commit a006a14
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/polars/polars-io/src/csv/write_impl.rs b/polars/polars-io/src/csv/write_impl.rs
@@ -1,6 +1,6 @@
 use arrow::temporal_conversions;
 use lexical_core::{FormattedSize, ToLexical};
-use memchr::memchr;
+use memchr::{memchr, memchr2};
 use polars_core::{prelude::*, series::SeriesIter, POOL};
 use polars_utils::contention_pool::LowContentionPool;
 use rayon::prelude::*;
@@ -23,7 +23,7 @@ fn fmt_and_escape_str(f: &mut Vec<u8>, v: &str, options: &SerializeOptions) -> s
             };
             return write!(f, "\"{}\"", replaced);
         }
-        let surround_with_quotes = memchr(options.delimiter, v.as_bytes()).is_some();
+        let surround_with_quotes = memchr2(options.delimiter, b'\n', v.as_bytes()).is_some();
 
         if surround_with_quotes {
             write!(f, "\"{}\"", v)

diff --git a/py-polars/tests/io/test_csv.py b/py-polars/tests/io/test_csv.py
@@ -582,3 +582,12 @@ def test_different_eol_char() -> None:
     assert pl.read_csv(csv.encode(), eol_char=";", has_header=False).frame_equal(
         expected
     )
+
+
+def test_csv_write_escape_newlines() -> None:
+    df = pl.DataFrame(dict(escape=["n\nn"]))
+    f = io.BytesIO()
+    df.write_csv(f)
+    f.seek(0)
+    read_df = pl.read_csv(f)
+    assert df.frame_equal(read_df)