examples/select_csv_columns.py

#!/usr/bin/env python3
"""Select named CSV fields without breaking quoted commas or multiline values.

Example (no source overwrite):
 python3 examples/select_csv_columns.py input.csv output.csv id,notes

This is a local example, not an embedded plugin execution capability.
"""
import argparse
import csv
import os
import tempfile
from pathlib import Path


def select_columns(source: Path, destination: Path, columns: list[str]) -> int:
    if not columns or len(set(columns))!=len(columns):
        raise ValueError("Specify unique column names")
    if source.resolve()==destination.resolve():
        raise ValueError("Input and output must differ; no in-place overwrite")
    if destination.exists():
        raise FileExistsError("Output exists; select a new path instead of overwriting")
    count=0
    temporary=None
    with source.open('r',newline='',encoding='utf-8-sig') as src:
        reader=csv.DictReader(src)
        if not reader.fieldnames:
            raise ValueError('CSV header missing')
        if len(reader.fieldnames)!=len(set(reader.fieldnames)):
            raise ValueError('Duplicate source column names')
        missing=[name for name in columns if name not in reader.fieldnames]
        if missing:raise ValueError('Columns not present: '+', '.join(missing))
        # Create an isolated staging file on the destination filesystem.
        fd,path=tempfile.mkstemp(prefix=f'.{destination.name}.',suffix='.tmp',dir=destination.parent)
        temporary=Path(path)
        try:
            with os.fdopen(fd,'w',newline='',encoding='utf-8') as dst:
                writer=csv.DictWriter(dst,fieldnames=columns,extrasaction='ignore')
                writer.writeheader()
                for row in reader:
                    if None in row:
                        raise ValueError(f'Row {reader.line_num} has more values than header')
                    if any(row[name] is None for name in reader.fieldnames):
                        raise ValueError(f'Row {reader.line_num} is missing values')
                    writer.writerow({name:row[name] for name in columns})
                    count+=1
                dst.flush()
                os.fsync(dst.fileno())
            # Link only when output does not exist. No accidental overwrite, even
            # when the destination appears between the initial check and publish.
            os.link(temporary,destination)
        finally:
            temporary.unlink(missing_ok=True)
    return count


def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('input',type=Path)
    p.add_argument('output',type=Path)
    p.add_argument('columns',help='comma-separated column names')
    a=p.parse_args()
    n=select_columns(a.input,a.output,[x.strip() for x in a.columns.split(',')])
    print(f'Wrote {n} records to {a.output}')

if __name__=='__main__':main()