“Know how to solve every problem that has been solved.” “What I cannot create, I do not understand.” — Richard Feynman
python 89 lines · 3.2 KB
Raw ↗
"""Deterministic test data for the CSV engine.

Writes two files next to this script:

  bench.csv  -- 1,000,000 data rows + header, ~70 MB, CRLF line endings.
                Columns: id(int), category(str), name(str), value(double),
                count(int), notes(str). The category column includes one value
                with an embedded comma (always quoted); the name column
                sprinkles RFC-4180 escapes: doubled quotes every 97th row and
                an embedded newline every 1009th row. Regenerate any time --
                bench.csv is gitignored; THIS FILE is its provenance.

  edge.csv   -- tiny RFC-4180 torture file. Its exact parsed content is
                hard-coded in csv_test_common.hpp; every step must reproduce
                it byte-for-byte before it is allowed to print a benchmark.

Also prints the reference aggregates for the step-5 task (sum of `value`,
per-category counts) so any run can be eyeballed against the C++ output.
"""
import random
from pathlib import Path

HERE = Path(__file__).parent
N_ROWS = 1_000_000
rng = random.Random(42)

CATEGORIES = ["alpha", "beta", "gamma", "delta", "epsilon",
              "widgets, large", "zeta", "eta"]


def quote_if_needed(field: str) -> str:
    if any(c in field for c in ',"\r\n'):
        return '"' + field.replace('"', '""') + '"'
    return field


def gen_bench():
    total_value = 0.0
    cat_counts = {c: 0 for c in CATEGORIES}
    with open(HERE / "bench.csv", "wb") as f:
        f.write(b"id,category,name,value,count,notes\r\n")
        for i in range(N_ROWS):
            cat = CATEGORIES[rng.randrange(len(CATEGORIES))]
            cat_counts[cat] += 1
            name = f"item-{i:07d}"
            if i % 97 == 0:
                name = f'item said "{i}" loudly'
            elif i % 1009 == 0:
                name = f"item-{i}\nsecond line"
            value = round(rng.uniform(0.0, 10000.0), 6)
            total_value += value
            count = rng.randrange(0, 100000)
            notes = "" if i % 5 == 0 else f"n{i % 1000}"
            row = ",".join([
                str(i),
                quote_if_needed(cat),
                quote_if_needed(name),
                f"{value:.6f}",
                str(count),
                quote_if_needed(notes),
            ])
            f.write(row.encode() + b"\r\n")
    size = (HERE / "bench.csv").stat().st_size
    print(f"bench.csv: {N_ROWS} rows, {size/1e6:.1f} MB")
    print(f"  reference sum(value)  = {total_value!r}")
    for c in CATEGORIES:
        print(f"  reference count[{c!r}] = {cat_counts[c]}")


def gen_edge():
    # Exact bytes matter: CRLF terminators, an embedded CRLF inside quotes,
    # doubled quotes, empty quoted + unquoted fields, and NO trailing newline
    # after the final record (RFC 4180 makes it optional).
    records = (
        b"id,name,notes,value\r\n"
        b"1,plain,simple,3.14\r\n"
        b'2,"quoted, with comma","line1\r\nline2",2.71\r\n'
        b'3,"she said ""hi""",bare,1\r\n'
        b'4,,"",0\r\n'
        b'5," leading and trailing ",x,42\r\n'
        b"6,end,no-trailing-newline,7"
    )
    (HERE / "edge.csv").write_bytes(records)
    print(f"edge.csv: {len(records)} bytes, 7 records, no trailing newline")


if __name__ == "__main__":
    gen_edge()
    gen_bench()