beam: RFC4180 CSV parser+writer with roundtrip tests

wizard-work/beam/csv_rfc4180.py

@@ -0,0 +1,104 @@
+"""RFC 4180-ish CSV parser + writer with correct handling of quoted fields."""
+from __future__ import annotations
+
+
+def parse_csv(text, delimiter=",", quotechar='"'):
+    if text.startswith("\ufeff"):
+        text = text[1:]
+    rows = []
+    field = []
+    row = []
+    i = 0
+    n = len(text)
+    in_quotes = False
+    at_field_start = True
+
+    def flush_field():
+        nonlocal field, row, at_field_start
+        row.append("".join(field))
+        field = []
+        at_field_start = True
+
+    def flush_row():
+        nonlocal row, at_field_start
+        rows.append(row)
+        row = []
+        at_field_start = True
+
+    if text == "":
+        return rows
+    while i < n:
+        c = text[i]
+        if in_quotes:
+            if c == quotechar:
+                if i + 1 < n and text[i + 1] == quotechar:
+                    field.append(quotechar)
+                    i += 2
+                else:
+                    in_quotes = False
+                    i += 1
+            else:
+                field.append(c)
+                i += 1
+        else:
+            if c == quotechar and at_field_start:
+                in_quotes = True
+                at_field_start = False
+                i += 1
+            elif c == delimiter:
+                flush_field()
+                i += 1
+            elif c == "\r":
+                flush_field()
+                flush_row()
+                if i + 1 < n and text[i + 1] == "\n":
+                    i += 2
+                else:
+                    i += 1
+                at_field_start = True
+            elif c == "\n":
+                flush_field()
+                flush_row()
+                i += 1
+                at_field_start = True
+            else:
+                field.append(c)
+                at_field_start = False
+                i += 1
+    if field or row or not at_field_start:
+        if not at_field_start or field:
+            flush_field()
+        if row:
+            flush_row()
+    return rows
+
+
+def _needs_quoting(field, delimiter, quotechar, lineterminator):
+    if quotechar in field or delimiter in field:
+        return True
+    if "\n" in field or "\r" in field:
+        return True
+    if field != field.strip(" "):
+        return True
+    return False
+
+
+def write(rows, delimiter=",", quotechar='"', lineterminator="\r\n"):
+    out = []
+    for row in rows:
+        fields = []
+        for f in row:
+            f = "" if f is None else str(f)
+            if _needs_quoting(f, delimiter, quotechar, lineterminator):
+                f = quotechar + f.replace(quotechar, quotechar + quotechar) + quotechar
+            fields.append(f)
+        out.append(delimiter.join(fields))
+    return lineterminator.join(out) + (lineterminator if rows else "")
+
+
+def dumps(rows, **kw):
+    return write(rows, **kw)
+
+
+def loads(text, **kw):
+    return parse_csv(text, **kw)

wizard-work/beam/test_csv.py

@@ -0,0 +1,35 @@
+import sys
+sys.path.insert(0, "/root/.hermes/wizard-outputs/beam")
+from csv_rfc4180 import parse_csv, write, loads, dumps
+
+def eq(a,b):
+    assert a==b, (a,b)
+
+# basic
+eq(parse_csv("a,b,c\n1,2,3\n"), [["a","b","c"],["1","2","3"]])
+# quoted embedded comma
+eq(parse_csv('a,"b,c",d\n'), [["a","b,c","d"]])
+# quoted newline
+eq(parse_csv('a,"b\nc",d\n'), [["a","b\nc","d"]])
+# doubled quote
+eq(parse_csv('a,"b""c",d\n'), [["a",'b"c',"d"]])
+# CRLF
+eq(parse_csv("a,b\r\n1,2\r\n"), [["a","b"],["1","2"]])
+# BOM
+eq(parse_csv("\ufeffa,b\n"), [["a","b"]])
+# empty input
+eq(parse_csv(""), [])
+# trailing empty line -> blank record ['']
+eq(parse_csv("a,b\n\n"), [["a","b"],[""]])
+# ragged
+eq(parse_csv("a,b,c\n1,2\n"), [["a","b","c"],["1","2"]])
+# write roundtrip
+data = [["name","note"],["Alice","has, comma"],["Bob","line1\nline2"],["C","quote\"here"]]
+s = write(data)
+print(repr(s))
+eq(parse_csv(s), data)
+# write quoting minimal
+eq(write([["a b","c"]]), 'a b,c\r\n')
+# empty field
+eq(parse_csv("a,,c\n"), [["a","","c"]])
+print("ALL TESTS PASSED")