a621b94512
Break the console summary out by year instead of collapsing every requested year into one set of totals.
378 lines
15 KiB
Python
378 lines
15 KiB
Python
"""
|
|
uc_salary_pipeline.py
|
|
=====================
|
|
All-in-one pipeline: download UC employee wage data from ucinvestments/
|
|
UC-wages (https://github.com/ucinvestments/UC-wages), classify every
|
|
employee into one of 8 analytical categories (see classifier.py), and
|
|
write a per-employee CSV.
|
|
|
|
USAGE
|
|
-----
|
|
# Download every campus for every available year and classify:
|
|
python uc_salary_pipeline.py
|
|
|
|
# A subset:
|
|
python uc_salary_pipeline.py --years 2023 2024 --campuses UCB UCLA
|
|
|
|
# Classify already-downloaded data without hitting the network:
|
|
python uc_salary_pipeline.py --skip-download
|
|
|
|
DATA ACQUISITION
|
|
----------------
|
|
Downloading is delegated to uc_wages_downloader.py, which fetches raw JSON
|
|
from ucinvestments/UC-wages into --data-dir (default uc_wages_data/),
|
|
skipping files already present. See that module's docstring for details.
|
|
|
|
CLASSIFICATION
|
|
---------------
|
|
Delegated entirely to classifier.py — see that module's docstring for the
|
|
category definitions. This script only downloads data, maps UC-wages'
|
|
JSON fields onto the (title, gross_pay) shape classifier.classify_title()
|
|
expects, and writes the result.
|
|
|
|
OUTPUT
|
|
------
|
|
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
|
|
|
|
Optionally, a top-earners CSV (--top-earners-output) with the --top-n
|
|
(default 10) highest gross_pay employees per category (also broken down
|
|
further by --group-by, e.g. per year):
|
|
category, category_label, rank, year, campus, title,
|
|
gross_pay, regular_pay, overtime_pay, other_pay, first_name, last_name
|
|
"""
|
|
|
|
import argparse
|
|
import csv
|
|
import heapq
|
|
import itertools
|
|
import json
|
|
import logging
|
|
import re
|
|
import statistics
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import classifier
|
|
import uc_wages_downloader as downloader
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# LOGGING
|
|
# ---------------------------------------------------------------------------
|
|
logging.basicConfig(
|
|
format="%(asctime)s [%(levelname)s] %(message)s",
|
|
level=logging.INFO,
|
|
handlers=[
|
|
logging.FileHandler("uc_salary_pipeline.log", encoding="utf-8"),
|
|
logging.StreamHandler(sys.stdout),
|
|
],
|
|
)
|
|
log = logging.getLogger(__name__)
|
|
|
|
OUTPUT_COLUMNS = [
|
|
"year", "campus", "title", "category", "category_label",
|
|
"gross_pay", "regular_pay", "overtime_pay", "other_pay",
|
|
"first_name", "last_name",
|
|
]
|
|
|
|
|
|
def _dollar(s) -> float:
|
|
try:
|
|
return float(re.sub(r"[$,\s]", "", str(s)))
|
|
except (ValueError, TypeError):
|
|
return 0.0
|
|
|
|
|
|
def parse_uc_wages_record(record: dict) -> dict | None:
|
|
"""Map one UC-wages JSON record onto the fields classify_title() needs."""
|
|
title = (record.get("title") or "").strip()
|
|
if not title:
|
|
return None
|
|
regular = _dollar(record.get("basepay") or 0)
|
|
overtime = _dollar(record.get("overtimepay") or 0)
|
|
other = _dollar(record.get("adjustpay") or 0)
|
|
gross = _dollar(record.get("grosspay") or 0)
|
|
if gross == 0:
|
|
gross = regular + overtime + other
|
|
return {
|
|
"title": title,
|
|
"gross_pay": gross,
|
|
"regular_pay": regular,
|
|
"overtime_pay": overtime,
|
|
"other_pay": other,
|
|
"first_name": (record.get("firstname") or "").strip(),
|
|
"last_name": (record.get("lastname") or "").strip(),
|
|
}
|
|
|
|
|
|
def _group_sort_key(key: tuple) -> tuple:
|
|
"""Sort a (group_dims..., category) key by its group values, then by
|
|
category in classifier.CATEGORIES order (not alphabetically)."""
|
|
*prefix, category = key
|
|
return (*prefix, classifier.CATEGORIES.index(category))
|
|
|
|
|
|
def process_file(path: Path, campus: str, year: int, writer: csv.writer,
|
|
senior_threshold: float, support_threshold: float,
|
|
agg: dict | None = None, group_dims: tuple = (),
|
|
top_earners: dict | None = None, top_n: int = 10,
|
|
tie_breaker: itertools.count | None = None,
|
|
console_stats: dict | None = None) -> int:
|
|
"""Classify every record in one UC-wages JSON file, writing a CSV row
|
|
per employee. Returns the number of rows written.
|
|
|
|
If `agg` is given, also accumulates each record's gross_pay under a
|
|
(group_dims..., category) key for later aggregate reporting.
|
|
|
|
If `top_earners` is given, also maintains a size-`top_n` min-heap of
|
|
the highest gross_pay employees per (group_dims..., category) key.
|
|
|
|
If `console_stats` is given, also accumulates each record's gross_pay
|
|
under its (year, category) for the end-of-run summary tables printed
|
|
to stdout (one table per year).
|
|
"""
|
|
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)
|
|
if parsed is None:
|
|
continue
|
|
category = classifier.classify_title(
|
|
parsed["title"],
|
|
gross_pay=parsed["gross_pay"],
|
|
senior_threshold=senior_threshold,
|
|
support_threshold=support_threshold,
|
|
)
|
|
writer.writerow([
|
|
year, campus, parsed["title"], category,
|
|
classifier.CATEGORY_LABELS[category],
|
|
round(parsed["gross_pay"], 2), round(parsed["regular_pay"], 2),
|
|
round(parsed["overtime_pay"], 2), round(parsed["other_pay"], 2),
|
|
parsed["first_name"], parsed["last_name"],
|
|
])
|
|
count += 1
|
|
if agg is not None:
|
|
agg.setdefault(group_prefix + (category,), []).append(parsed["gross_pay"])
|
|
if console_stats is not None:
|
|
console_stats.setdefault((year, category), []).append(parsed["gross_pay"])
|
|
if top_earners is not None:
|
|
entry = (parsed["gross_pay"], next(tie_breaker),
|
|
(year, campus, parsed["title"], parsed["gross_pay"],
|
|
parsed["regular_pay"], parsed["overtime_pay"],
|
|
parsed["other_pay"], parsed["first_name"], parsed["last_name"]))
|
|
heap = top_earners.setdefault(group_prefix + (category,), [])
|
|
if len(heap) < top_n:
|
|
heapq.heappush(heap, entry)
|
|
else:
|
|
heapq.heappushpop(heap, entry)
|
|
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",
|
|
]
|
|
|
|
with open(outpath, "w", newline="", encoding="utf-8") as fh:
|
|
writer = csv.writer(fh)
|
|
writer.writerow(header)
|
|
for key in sorted(agg, key=_group_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)
|
|
|
|
|
|
def write_top_earners_csv(top_earners: dict, outpath: Path) -> None:
|
|
"""Write the top-N earners per (group_dims..., category) key, ranked
|
|
highest gross_pay first within each group."""
|
|
header = ["category", "category_label", "rank", "year", "campus", "title",
|
|
"gross_pay", "regular_pay", "overtime_pay", "other_pay",
|
|
"first_name", "last_name"]
|
|
|
|
with open(outpath, "w", newline="", encoding="utf-8") as fh:
|
|
writer = csv.writer(fh)
|
|
writer.writerow(header)
|
|
for key in sorted(top_earners, key=_group_sort_key):
|
|
category = key[-1]
|
|
ranked = sorted(top_earners[key], key=lambda e: e[0], reverse=True)
|
|
for rank, (gross_pay, _, row) in enumerate(ranked, start=1):
|
|
year, campus, title, gross, regular, overtime, other, first, last = row
|
|
writer.writerow([
|
|
category, classifier.CATEGORY_LABELS[category], rank,
|
|
year, campus, title, round(gross, 2), round(regular, 2),
|
|
round(overtime, 2), round(other, 2), first, last,
|
|
])
|
|
log.info("Top earners CSV -> %s", outpath)
|
|
|
|
|
|
def _money(v: float) -> str:
|
|
if abs(v) >= 1_000_000:
|
|
return f"${v / 1_000_000:,.2f}M"
|
|
return f"${v:,.0f}"
|
|
|
|
|
|
def print_console_summary(console_stats: dict) -> None:
|
|
"""Print one formatted category-by-category summary table per year,
|
|
keyed off (year, category) -> [gross_pay, ...] entries in console_stats."""
|
|
label_w, col_w = 34, 16
|
|
header = (f"{'Category':<{label_w}}{'Employees':>{col_w}}"
|
|
f"{'Total Pay':>{col_w}}{'Median Pay':>{col_w}}")
|
|
rule = "=" * len(header)
|
|
|
|
years = sorted({year for year, _ in console_stats})
|
|
for year in years:
|
|
print(f"\n{rule}")
|
|
print(f"UC SALARY CLASSIFICATION SUMMARY -- {year}")
|
|
print(rule)
|
|
print(header)
|
|
print("-" * len(header))
|
|
|
|
total_count, total_pay = 0, 0.0
|
|
for category in classifier.CATEGORIES:
|
|
values = console_stats.get((year, category), [])
|
|
count, total = len(values), sum(values)
|
|
median = statistics.median(values) if values else 0.0
|
|
total_count += count
|
|
total_pay += total
|
|
label = classifier.CATEGORY_LABELS[category]
|
|
print(f"{label:<{label_w}}{count:>{col_w},}"
|
|
f"{_money(total):>{col_w}}{_money(median):>{col_w}}")
|
|
|
|
print("-" * len(header))
|
|
print(f"{'TOTAL':<{label_w}}{total_count:>{col_w},}{_money(total_pay):>{col_w}}")
|
|
print(rule)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Download and classify UC employee wage data into 8 categories.",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog=__doc__,
|
|
)
|
|
parser.add_argument("--years", nargs="+", type=int,
|
|
default=downloader.DEFAULT_YEARS,
|
|
help=f"Year(s) to process (default: 2010-"
|
|
f"{downloader.DEFAULT_YEARS[-1]}).")
|
|
parser.add_argument("--campuses", nargs="+",
|
|
choices=list(downloader.LOCATIONS.keys()),
|
|
default=list(downloader.LOCATIONS.keys()),
|
|
metavar="CAMPUS",
|
|
help="Campuses to process (default: all). Choices: " +
|
|
" ".join(downloader.LOCATIONS.keys()))
|
|
parser.add_argument("--data-dir", default="uc_wages_data",
|
|
help="Local cache directory for downloaded JSON "
|
|
"(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 and --top-earners-output, "
|
|
"alongside category (default: year). Pass both "
|
|
"for a year x campus x category breakdown, or "
|
|
"nothing to collapse everything to one group "
|
|
"per category.")
|
|
parser.add_argument("--top-earners-output",
|
|
help="If given, also write a CSV of the --top-n "
|
|
"highest gross_pay employees per category (see "
|
|
"--group-by).")
|
|
parser.add_argument("--top-n", type=int, default=10,
|
|
help="Number of top earners per group for "
|
|
"--top-earners-output (default: 10).")
|
|
parser.add_argument("--skip-download", action="store_true",
|
|
help="Classify whatever is already cached in "
|
|
"--data-dir without hitting the network.")
|
|
parser.add_argument("--quiet", action="store_true",
|
|
help="Suppress the formatted summary table printed "
|
|
"to stdout after processing.")
|
|
parser.add_argument("--senior-threshold", type=float,
|
|
default=classifier.SENIOR_STAFF_THRESHOLD,
|
|
help=f"Gross pay threshold for SENIOR_PROF_STAFF "
|
|
f"(default ${classifier.SENIOR_STAFF_THRESHOLD:,}).")
|
|
parser.add_argument("--support-threshold", type=float,
|
|
default=classifier.SUPPORT_STAFF_THRESHOLD,
|
|
help=f"Gross pay below this -> SUPPORT_STAFF "
|
|
f"(default ${classifier.SUPPORT_STAFF_THRESHOLD:,}).")
|
|
args = parser.parse_args()
|
|
|
|
data_dir = Path(args.data_dir)
|
|
|
|
if not args.skip_download:
|
|
downloader.download_all(args.years, args.campuses, data_dir)
|
|
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
|
|
top_earners = {} if args.top_earners_output else None
|
|
tie_breaker = itertools.count() if top_earners is not None else None
|
|
console_stats = {} if not args.quiet else None
|
|
|
|
output_path = Path(args.output)
|
|
total_rows = 0
|
|
with open(output_path, "w", newline="", encoding="utf-8") as out_fh:
|
|
writer = csv.writer(out_fh)
|
|
writer.writerow(OUTPUT_COLUMNS)
|
|
|
|
for year in sorted(args.years):
|
|
for campus in args.campuses:
|
|
path = downloader.out_path(campus, year, data_dir)
|
|
if not path.exists():
|
|
log.info("No data for %s %s, skipping.", campus, year)
|
|
continue
|
|
count = process_file(path, campus, year, writer,
|
|
args.senior_threshold, args.support_threshold,
|
|
agg=agg, group_dims=group_dims,
|
|
top_earners=top_earners, top_n=args.top_n,
|
|
tie_breaker=tie_breaker,
|
|
console_stats=console_stats)
|
|
total_rows += count
|
|
log.info("%s %s: classified %d records", campus, year, count)
|
|
|
|
if total_rows == 0:
|
|
log.error("No data processed.")
|
|
sys.exit(1)
|
|
|
|
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 top_earners is not None:
|
|
write_top_earners_csv(top_earners, Path(args.top_earners_output))
|
|
|
|
if console_stats is not None:
|
|
print_console_summary(console_stats)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|