mirror of
https://github.com/tgorordo/smithy.git
synced 2026-06-05 16:22:15 -07:00
115 lines
3.1 KiB
Python
115 lines
3.1 KiB
Python
# /// script
|
|
# requires-python = ">=3.13"
|
|
# dependencies = [
|
|
# "click>=8.4.1",
|
|
# "rich>=15.0.0",
|
|
# "polars>=1.40.1",
|
|
# "rustworkx>=0.17.1"
|
|
# ]
|
|
# ///
|
|
import sys, io
|
|
import click
|
|
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
from rich.panel import Panel
|
|
|
|
import polars as pl
|
|
from smithy import smith_set, irv_set
|
|
|
|
|
|
@click.command()
|
|
@click.argument("ballots", type=click.Path(exists=True, dir_okay=False))
|
|
@click.option(
|
|
"--try-resolve-irv",
|
|
is_flag=True,
|
|
help="Try to reduce or resolve the Smith set by running all-paths IRV on the set.",
|
|
)
|
|
@click.option(
|
|
"--show-ballots",
|
|
"-b",
|
|
is_flag=True,
|
|
help="Show relevant ballots (after selections).",
|
|
)
|
|
@click.option("--pretty", "-p", is_flag=True, help="Pretty-print output.")
|
|
def cli(ballots: str, try_resolve_irv=False, show_ballots=False, pretty=False) -> None:
|
|
"""
|
|
Compute the Smith set from a box of ranked-choice ballots -- .csv or .xls(x).
|
|
|
|
The Smith set is the minimal set of candidates which can beat all others pairwise
|
|
(simple ranking majority) - if there is a single winner in the set,
|
|
they are guaranteed the Condorcet i.e. Majority winner.
|
|
"""
|
|
|
|
console = Console()
|
|
|
|
try:
|
|
# Load ballots
|
|
if ballots.endswith(".csv"):
|
|
df = pl.read_csv(ballots)
|
|
|
|
elif ballots.endswith((".xlsx", ".xls")):
|
|
df = pl.read_excel(ballots)
|
|
|
|
else:
|
|
console.print(
|
|
"[bold red]Unsupported file type.[/bold red]\nUse CSV or Excel."
|
|
)
|
|
raise SystemExit(1)
|
|
|
|
# Normalize numerical dataframe entries
|
|
df = df.with_columns(
|
|
[
|
|
pl.col(c)
|
|
.cast(pl.String)
|
|
.str.strip_chars()
|
|
.cast(pl.Int64, strict=False)
|
|
.fill_null(0)
|
|
for c in df.columns
|
|
]
|
|
)
|
|
|
|
# Compute Smith set
|
|
smiths = smith_set(df)
|
|
if len(smiths) > 1 and try_resolve_irv:
|
|
irv_ballots = df.select(smiths)
|
|
smiths = irv_set(irv_ballots)
|
|
|
|
if show_ballots and pretty:
|
|
preview = Table(title="Ballot Box")
|
|
|
|
for col in df.columns:
|
|
preview.add_column(col)
|
|
|
|
for row in df.head(5).iter_rows():
|
|
preview.add_row(*map(str, row))
|
|
|
|
console.print(preview)
|
|
console.print()
|
|
elif show_ballots and not pretty:
|
|
buf = io.StringIO()
|
|
df.write_csv(buf)
|
|
console.print(buf.getvalue())
|
|
console.print("---")
|
|
|
|
if pretty:
|
|
console.print(
|
|
Panel.fit(
|
|
"\n".join(f"• {c}" for c in smiths),
|
|
title="Resulting IRV-resolved Smith Set"
|
|
if (try_resolve_irv)
|
|
else "Resulting Smith Set",
|
|
border_style="green",
|
|
)
|
|
)
|
|
else:
|
|
console.print(", ".join(f"{c}" for c in smiths))
|
|
|
|
except Exception as e:
|
|
console.print(f"[bold red]Error:[/bold red] {e}")
|
|
|
|
raise SystemExit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli()
|