tally.py

import argparse
import fractions
from collections import defaultdict, Counter
from itertools import chain

from tabulate import tabulate


def format_pos(position: int, primary_count: int) -> str:
    if primary_count:
        if position <= primary_count:
            return f"#{position}"
        return f"(#{position - primary_count})"
    return f"#{position}"


def format_fraction(frac: fractions.Fraction) -> str:
    return f"{frac.numerator}/{frac.denominator}"


def format_table_row(
    pos: int,
    name: str,
    name_votes: list[fractions.Fraction],
    extant_fractions: list[fractions.Fraction],
    *,
    primary_count: int,
):
    score = sum(name_votes)
    vote_ctr = Counter(name_votes)
    return (
        format_pos(pos, primary_count),
        name,
        float(score),
        str(score),
        *(vote_ctr[v] for v in extant_fractions),
    )


def format_table(
    votes: dict[str, list[fractions.Fraction]], primary_count: int
) -> None:
    extant_fractions = sorted(set(chain(*votes.values())), reverse=True)
    table = [
        format_table_row(
            pos,
            name,
            name_votes,
            primary_count=primary_count,
            extant_fractions=extant_fractions,
        )
        for pos, (name, name_votes) in enumerate(
            sorted(votes.items(), key=lambda p: sum(p[1]), reverse=True),
            1,
        )
    ]
    headers = (
        "Position",
        "Name",
        "Total",
        "Fraction",
        *(f"Σ {format_fraction(f)}" for f in extant_fractions),
    )
    print(tabulate(table, headers=headers))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("file")
    ap.add_argument("-c", "--count", type=int, required=True)
    ap.add_argument("-r", "--reverse", action="store_true", default=False)
    args = ap.parse_args()
    primary_count = args.count
    if primary_count < 1:
        raise ValueError("primary count must be >= 1")

    votes = defaultdict(list)

    with open(args.file) as infp:
        for line in infp:
            line = line.strip()
            if not line:
                continue
            if line.startswith("#"):
                continue
            name, num = line.split(None, 1)
            if args.reverse:
                name, num = num, name
            num = int(num)
            if not (1 <= num <= primary_count):
                raise ValueError(f"invalid line: {line}")
            frac = fractions.Fraction(1, num)
            votes[name.title()].append(frac)
    format_table(votes, primary_count=primary_count)


if __name__ == "__main__":
    main()