diff --git a/uc_salary_pipeline.py b/uc_salary_pipeline.py index 95e9376..ef2fa2f 100644 --- a/uc_salary_pipeline.py +++ b/uc_salary_pipeline.py @@ -32,9 +32,14 @@ expects, and writes the result. OUTPUT ------ -One row per employee per year: +Per-employee CSV (--output), one row per employee per year: year, campus, title, category, category_label, gross_pay, regular_pay, overtime_pay, other_pay, first_name, last_name + +Optionally, an aggregate CSV (--aggregate-output) with one row per +category (broken down further by --group-by year/campus): + [year] [campus] category, category_label, + employee_count, total_gross_pay, mean_gross_pay, median_gross_pay """ import argparse @@ -42,6 +47,7 @@ import csv import json import logging import re +import statistics import sys from pathlib import Path @@ -98,12 +104,21 @@ def parse_uc_wages_record(record: dict) -> dict | None: def process_file(path: Path, campus: str, year: int, writer: csv.writer, - senior_threshold: float, support_threshold: float) -> int: + senior_threshold: float, support_threshold: float, + agg: dict | None = None, group_dims: tuple = ()) -> int: """Classify every record in one UC-wages JSON file, writing a CSV row - per employee. Returns the number of rows written.""" + per employee. Returns the number of rows written. + + If `agg` is given, also accumulates each record's gross_pay under an + (group_dims..., category) key for later aggregate reporting. + """ with open(path, encoding="utf-8") as fh: data = json.load(fh) + group_prefix = tuple( + year if dim == "year" else campus for dim in group_dims + ) + count = 0 for record in data.get("records", []): parsed = parse_uc_wages_record(record) @@ -123,9 +138,38 @@ def process_file(path: Path, campus: str, year: int, writer: csv.writer, parsed["first_name"], parsed["last_name"], ]) count += 1 + if agg is not None: + agg.setdefault(group_prefix + (category,), []).append(parsed["gross_pay"]) return count +def write_aggregate_csv(agg: dict, group_dims: tuple, outpath: Path) -> None: + """Write one row per (group_dims..., category) with employee count, + total/mean/median gross pay.""" + header = list(group_dims) + [ + "category", "category_label", + "employee_count", "total_gross_pay", "mean_gross_pay", "median_gross_pay", + ] + + def sort_key(key): + *prefix, category = key + return (*prefix, classifier.CATEGORIES.index(category)) + + with open(outpath, "w", newline="", encoding="utf-8") as fh: + writer = csv.writer(fh) + writer.writerow(header) + for key in sorted(agg, key=sort_key): + *prefix, category = key + values = agg[key] + total = sum(values) + writer.writerow([ + *prefix, category, classifier.CATEGORY_LABELS[category], + len(values), round(total, 2), round(total / len(values), 2), + round(statistics.median(values), 2), + ]) + log.info("Aggregate CSV -> %s", outpath) + + # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- @@ -151,6 +195,17 @@ def main(): "(default: uc_wages_data/).") parser.add_argument("--output", default="uc_salary_classified.csv", help="Output CSV path (default: uc_salary_classified.csv).") + parser.add_argument("--aggregate-output", + help="If given, also write an aggregate CSV here with " + "employee count, total/mean/median gross pay per " + "category (see --group-by).") + parser.add_argument("--group-by", nargs="*", choices=["year", "campus"], + default=["year"], + help="Additional breakdown dimensions for " + "--aggregate-output, alongside category " + "(default: year). Pass both for a year x campus " + "x category breakdown, or nothing to collapse " + "everything to one row per category.") parser.add_argument("--skip-download", action="store_true", help="Classify whatever is already cached in " "--data-dir without hitting the network.") @@ -171,6 +226,9 @@ def main(): else: log.info("--skip-download set, classifying cached data only.") + group_dims = tuple(dim for dim in ("year", "campus") if dim in args.group_by) + agg = {} if args.aggregate_output else None + output_path = Path(args.output) total_rows = 0 with open(output_path, "w", newline="", encoding="utf-8") as out_fh: @@ -184,7 +242,8 @@ def main(): log.info("No data for %s %s, skipping.", campus, year) continue count = process_file(path, campus, year, writer, - args.senior_threshold, args.support_threshold) + args.senior_threshold, args.support_threshold, + agg=agg, group_dims=group_dims) total_rows += count log.info("%s %s: classified %d records", campus, year, count) @@ -194,6 +253,9 @@ def main(): log.info("Total rows written: %d -> %s", total_rows, output_path) + if agg is not None: + write_aggregate_csv(agg, group_dims, Path(args.aggregate_output)) + if __name__ == "__main__": main()