Add places to a ZIP-code CSV

import argparse
import csv
import io
import re
import sys
from pathlib import Path

from parseapi import ParseAPI, ParseAPIError

SAMPLE = """id,zip,note
1,02108,Boston
2,94107,San Francisco
3,02108-1234,Same five-digit lookup
4,2108,Missing digit
5,,Missing ZIP
6,00000,Unknown ZIP
"""
COLUMNS = ["parse_postal", "parse_city", "parse_state", "parse_timezone",
           "parse_status", "parse_error", "parse_http_status"]


# example:read
def read_rows(stream):
    reader = csv.reader(stream, strict=True)
    fields = next(reader, None)
    if (not fields or "zip" not in fields or any(not field.strip() for field in fields)
            or len(fields) != len(set(fields))):
        raise ValueError("Use distinct, nonempty headers, including a column named zip.")
    if set(fields) & set(COLUMNS):
        raise ValueError("The input already contains output columns. Use the original CSV.")
    rows = []
    for cells in reader:
        if len(cells) != len(fields):
            raise ValueError("Every row must contain the same number of cells as the header.")
        rows.append(dict(zip(fields, cells)))
        if len(rows) > 200:
            raise ValueError("Use a CSV with 1 to 200 data rows for this example.")
    if not rows:
        raise ValueError("The CSV needs at least one data row.")
    return fields, rows
# /example:read


def outcome(status, error="", http_status=""):
    return {**dict.fromkeys(COLUMNS, ""), "parse_status": status,
            "parse_error": error, "parse_http_status": http_status}


def matched(data, postal):
    if (not isinstance(data, dict) or data.get("postal") != postal
            or data.get("country") != "US"):
        return outcome("error", "unexpected_response")
    names = ("city", "state", "timezone")
    if any(name not in data or (data[name] is not None and
           (not isinstance(data[name], str) or not data[name].strip())) for name in names):
        return outcome("error", "unexpected_response")
    # example:result
    return {**outcome("matched", http_status=200), "parse_postal": postal,
            **{f"parse_{name}": data[name] or "" for name in names}}
    # /example:result


def enrich(rows, parse):
    results = []
    cache = {}
    stopped = False
    for row in rows:
        value = row["zip"].strip()
        if not value:
            result = outcome("missing_zip")
        elif not re.fullmatch(r"[0-9]{5}(?:-[0-9]{4})?", value):
            result = outcome("invalid_zip_format")
        else:
            postal = value[:5]
            if postal in cache:
                result = cache[postal]
            elif stopped:
                result = outcome("not_attempted")
            else:
                try:
                    # example:request
                    result = matched(parse.postal(postal, country="US"), postal)
                    # /example:request
                except ParseAPIError as error:
                    status = "not_found" if error.status == 404 else "error"
                    result = outcome(status, error.code, error.status)
                    if error.status in (401, 403, 429):
                        stopped = True
                except Exception:
                    result = outcome("error", "request_failed")
                # Reuse completed answers and failures for duplicate ZIPs in this run.
                cache[postal] = result
        results.append({**row, **result})
    return results


def write_rows(stream, fields, rows):
    writer = csv.DictWriter(stream, fieldnames=fields + COLUMNS)
    writer.writeheader()
    writer.writerows(rows)


def main():
    parser = argparse.ArgumentParser(description="Add city, state, and timezone to a US ZIP CSV.")
    parser.add_argument("input", nargs="?", help="UTF-8 CSV with a zip column; defaults to a sample")
    parser.add_argument("--output", help="New CSV filename; defaults to standard output")
    args = parser.parse_args()
    try:
        if args.output and (Path(args.output).exists() or Path(args.output).is_symlink()):
            raise ValueError("The output path already exists. Choose a new filename.")
        if args.input:
            path = Path(args.input)
            if path.stat().st_size > 65536:
                raise ValueError("Use a CSV no larger than 64 KiB.")
            with path.open(encoding="utf-8-sig", newline="") as stream:
                fields, rows = read_rows(stream)
        else:
            fields, rows = read_rows(io.StringIO(SAMPLE))
        # Reads PARSEAPI_KEY. Reuse one client, with no automatic retries.
        with ParseAPI(timeout=10.0, retries=0) as parse:
            results = enrich(rows, parse)
        if args.output:
            # Exclusive creation also protects against a file appearing during the lookups.
            with Path(args.output).open("x", encoding="utf-8", newline="") as stream:
                write_rows(stream, fields, results)
        else:
            write_rows(sys.stdout, fields, results)
    except (ValueError, OSError, csv.Error) as error:
        print(f"Could not complete: {error}", file=sys.stderr)
        return 1
    review = sum(row["parse_status"] != "matched" for row in results)
    print(f"{len(results)} rows written; {review} need review.", file=sys.stderr)
    return 2 if review else 0


if __name__ == "__main__":
    sys.exit(main())