Merge CSV Shards in Numeric Order with Python: Headers and Failure-Safe Output

from pathlib import Path

print(sorted(path.name for path in Path("parts").glob("part-*.csv")))
# ['part-1.csv', 'part-10.csv', 'part-2.csv']

python3 --version
mkdir csv-merge-demo
cd csv-merge-demo
mkdir parts
python3 - <<'PY_DATA'
import csv
from pathlib import Path

for number, note in [(10, "last"), (1, "first line\nsecond line"),
                     (2, "comma, inside")]:
    path = Path("parts") / f"part-{number}.csv"
    with path.open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["id", "note"])
        writer.writerow([f"{number:03d}", note])
PY_DATA

"""Merge numbered CSV shards in numeric filename order. Python 3.9+."""
import argparse
import csv
import os
from pathlib import Path
import re
import sys
import tempfile

NAME = re.compile(r"part-([0-9]+)\.csv")


def ordered_parts(directory):
    parts = []
    seen = set()
    for path in directory.glob("part-*.csv"):
        match = NAME.fullmatch(path.name)
        if match is None or path.is_symlink() or not path.is_file():
            raise ValueError(f"invalid shard: {path.name}")
        number = int(match[1])
        if number < 1 or number in seen:
            raise ValueError(f"nonpositive or duplicate shard number: {path.name}")
        seen.add(number)
        parts.append((number, path))
    if not parts:
        raise ValueError("no part-<positive integer>.csv files")
    return sorted(parts, key=lambda item: item[0])


def merge_csv(directory, output):
    directory = Path(directory).expanduser().resolve(strict=True)
    if not directory.is_dir():
        raise ValueError("input must be a directory")
    output = Path(output).expanduser().absolute()
    output = output.parent.resolve(strict=True) / output.name
    if output.parent == directory or output.is_symlink():
        raise ValueError("output must be outside input directory and not a symlink")
    parts = ordered_parts(directory)
    temporary = None
    count = 0
    expected = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="w", encoding="utf-8", newline="", dir=output.parent,
            prefix=f".{output.name}.", suffix=".tmp", delete=False,
        ) as dest:
            temporary = Path(dest.name)
            writer = csv.writer(dest, lineterminator="\n")
            for number, path in parts:
                with path.open(encoding="utf-8-sig", newline="") as source:
                    reader = csv.reader(source, strict=True)
                    try:
                        header = next(reader, None)
                        if not header or any(not field.strip() for field in header):
                            raise ValueError(f"{path.name}: missing or empty header")
                        if len(set(header)) != len(header):
                            raise ValueError(f"{path.name}: duplicate column names")
                        if expected is None:
                            expected = header
                            writer.writerow(header)
                        elif header != expected:
                            raise ValueError(f"{path.name}: header mismatch")
                        for row in reader:
                            if len(row) != len(expected):
                                raise ValueError(
                                    f"{path.name}: physical line {reader.line_num}: "
                                    f"expected {len(expected)} fields, got {len(row)}"
                                )
                            writer.writerow(row)
                            count += 1
                    except csv.Error as exc:
                        raise ValueError(
                            f"{path.name}: physical line {reader.line_num}: {exc}"
                        ) from exc
        os.replace(temporary, output)
    finally:
        if temporary is not None:
            temporary.unlink(missing_ok=True)
    return len(parts), count, [number for number, path in parts]


def main():
    parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False)
    parser.add_argument("directory", type=Path, help="directory of completed shards")
    parser.add_argument("output", type=Path, help="destination outside that directory")
    args = parser.parse_args()
    try:
        files, records, order = merge_csv(args.directory, args.output)
    except (OSError, ValueError, RuntimeError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1
    print(f"OK: files={files} records={records} order={','.join(map(str, order))}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

python3 merge_csv.py --help
python3 merge_csv.py parts merged.csv

OK: files=3 records=3 order=1,2,10

python3 - <<'PY_CHECK'
import csv

with open("merged.csv", encoding="utf-8", newline="") as handle:
    rows = list(csv.reader(handle))
assert rows == [
    ["id", "note"],
    ["001", "first line\nsecond line"],
    ["002", "comma, inside"],
    ["010", "last"],
]
print("PASS: numeric order, one header, and multiline values")
PY_CHECK

python3 - <<'PY_FAIL'
from pathlib import Path
import subprocess
import sys

output = Path("merged.csv")
before = output.read_bytes()
bad = Path("parts/part-10.csv")
original = bad.read_bytes()
try:
    bad.write_text("note,id\nwrong,010\n", encoding="utf-8")
    result = subprocess.run(
        [sys.executable, "merge_csv.py", "parts", "merged.csv"],
        capture_output=True, text=True,
    )
    assert result.returncode == 1, result
    assert result.stdout == "", result.stdout
    assert "header mismatch" in result.stderr, result.stderr
    assert output.read_bytes() == before
    assert not list(Path(".").glob(".merged.csv.*.tmp"))
    print("PASS: header mismatch returns 1; previous output unchanged")
finally:
    bad.write_bytes(original)
PY_FAIL

cp parts/part-1.csv parts/part-01.csv
python3 merge_csv.py parts merged.csv
printf 'exit=%s\n' "$?"
rm parts/part-01.csv