216 lines
7.9 KiB
Python
216 lines
7.9 KiB
Python
import json
|
|
import random
|
|
from argparse import ArgumentParser
|
|
from pathlib import Path
|
|
from typing import DefaultDict
|
|
|
|
QUEENS_FILE = "queens.txt"
|
|
PICKS_FILE_GLOB = "picks-*.txt"
|
|
ELIMINATIONS_FILE = "eliminations.txt"
|
|
ROSTERS_FILE = "rosters.json"
|
|
|
|
|
|
class Draft:
|
|
"""A draft for a season of RuPaul's Drag Race
|
|
|
|
This class represents a draft for a season. It is initialized with a path to a file
|
|
containing the queens for the season, and a glob pattern for the files containing the
|
|
picks from each player in the for of `picks-playername.txt`. This will allow the draft
|
|
to be run.
|
|
|
|
Alternatively, it can be initialized with a `roster.json` file which will load a pre-run
|
|
draft.
|
|
|
|
The draft itself is run by calling the `run` method. This will chose a random order for the players
|
|
in the draft and print it out. It will then then go through each player in order and pick the top ranked queen that
|
|
is still available. The draft will then print out the final rosters for each player.
|
|
|
|
The validate method should be called before running the draft to ensure that the queens and picks files are valid.
|
|
In particular:
|
|
* All the queens in the picks files should be in the queens file
|
|
* There should be less players than the total number of queens
|
|
* A player should not have picked the same queen twice
|
|
"""
|
|
|
|
def __init__(
|
|
self, queens_file: Path, picks_files: list[Path], roster_file: Path|None = None
|
|
):
|
|
self.queens = self.load_queens(queens_file)
|
|
self.players = self.load_players(picks_files)
|
|
if roster_file and roster_file.exists():
|
|
self.load_rosters(roster_file)
|
|
else:
|
|
self.rosters: dict[str, list[str]] = DefaultDict(list)
|
|
|
|
def load_queens(self, queens_file: Path) -> set[str]:
|
|
"""Load set of queens from the given file."""
|
|
return set(queens_file.read_text().strip().split("\n"))
|
|
|
|
def load_players(self, picks_files: list[Path]) -> dict[str, list[str]]:
|
|
"""Load the picks for each player from the given files."""
|
|
players: dict[str, list[str]] = {}
|
|
for picks_file in picks_files:
|
|
player = picks_file.stem.split("-")[1]
|
|
picks = picks_file.read_text().strip().split("\n")
|
|
if len(picks) != len(self.queens):
|
|
print(f"Player {player} has not picked all the queens")
|
|
|
|
picks_set = set(picks)
|
|
assert len(picks_set) == len(picks), f"Player {player} has duplicate picks"
|
|
unknown_queens = picks_set - self.queens
|
|
assert len(unknown_queens) == 0, f"Player {player} has invalid picks: {unknown_queens}"
|
|
|
|
players[player] = picks
|
|
|
|
return players
|
|
|
|
def load_rosters(self, roster_file: Path):
|
|
"""Load the rosters from the given file."""
|
|
self.rosters = json.loads(roster_file.read_text())
|
|
|
|
def validate(self):
|
|
"""Validate the draft before running."""
|
|
assert len(self.players) < len(
|
|
self.queens
|
|
), "There are too many players for the number of queens"
|
|
|
|
for player, picks in self.players.items():
|
|
assert len(picks) == len(
|
|
self.queens
|
|
), f"Player {player} has not picked all the queens"
|
|
assert len(set(picks)) == len(picks), f"Player {player} has duplicate picks"
|
|
assert set(picks) - self.queens, f"Player {player} has invalid picks"
|
|
|
|
def run(self):
|
|
"""Run the draft."""
|
|
self.order = list(self.players)
|
|
random.shuffle(self.order)
|
|
print("Draft order:")
|
|
for i, player in enumerate(self.order):
|
|
print(f"{i+1}. {player}")
|
|
|
|
remaining_queens = self.queens.copy()
|
|
|
|
round = 1
|
|
while len(remaining_queens) >= len(self.order):
|
|
print(f"Round {round}")
|
|
for player in self.order:
|
|
for queen in self.players[player]:
|
|
if queen in remaining_queens:
|
|
print(f"{player}: picks {queen}")
|
|
self.rosters[player].append(queen)
|
|
remaining_queens.remove(queen)
|
|
break
|
|
|
|
print("Draft complete")
|
|
self.print_rosters()
|
|
|
|
def print_rosters(self):
|
|
"""Print the final rosters for each player."""
|
|
for player, roster in self.rosters.items():
|
|
print(f"{player}: {', '.join(roster)}")
|
|
|
|
def save(self, roster_file: Path):
|
|
"""Save the rosters to the given file."""
|
|
roster_file.write_text(json.dumps(self.rosters))
|
|
|
|
def draft_completed(self):
|
|
"""Check if the draft has been completed."""
|
|
return len(self.rosters) == len(self.players)
|
|
|
|
def score_rosters(self, eliminations: list[str]) -> dict[str, int]:
|
|
"""Score each roster based on the queens that are still in the competition.
|
|
|
|
Each player should get a point for every queen that is still in the competition for each round of eliminations.
|
|
For example, if a player has 2 queens in the competition for a given round, they would get 2 points for that round.
|
|
If one of those queens was eliminated in the next round, they would only get 1 point for that round for a total
|
|
of 3 points.
|
|
|
|
The function should return a dictionary mapping each player to their total score as well as printing out the scores
|
|
at the end of each round."""
|
|
|
|
scores = {player: 0 for player in self.rosters}
|
|
remaining_queens = self.queens.copy()
|
|
for round, eliminated_queen in enumerate(eliminations):
|
|
remaining_queens.remove(eliminated_queen)
|
|
for player, roster in self.rosters.items():
|
|
scores[player] += len(set(roster) & remaining_queens)
|
|
|
|
print(f"Round {round+1} scores:")
|
|
for player, score in scores.items():
|
|
print(f"{player}: {score}")
|
|
|
|
return scores
|
|
|
|
|
|
def eliminate_queen(queens_file: Path, eliminations_file: Path, queen: str) -> list[str]|None:
|
|
"""Eliminate a queen from the list of queens."""
|
|
queens = set(queens_file.read_text().strip().split("\n"))
|
|
if queen not in queens:
|
|
print(f"{queen} is not a valid queen")
|
|
return None
|
|
|
|
eliminations = eliminations_file.read_text().strip().split("\n")
|
|
if queen in eliminations:
|
|
print(f"{queen} has already been eliminated")
|
|
return None
|
|
|
|
eliminations.append(queen)
|
|
eliminations_file.write_text("\n".join(eliminations))
|
|
|
|
return eliminations
|
|
|
|
|
|
def main():
|
|
parser = ArgumentParser()
|
|
parser.add_argument(
|
|
"season_dir",
|
|
type=Path,
|
|
default=Path("."),
|
|
help="Directory containing season data",
|
|
)
|
|
parser.add_argument(
|
|
"--eliminate", type=str, help="Eliminate a queen", metavar="QUEEN"
|
|
)
|
|
parser.add_argument("--draft", action="store_true", help="Run the initial draft")
|
|
parser.add_argument("--print", action="store_true", help="Print the current status")
|
|
args = parser.parse_args()
|
|
|
|
# Create a Draft object for the given season directory
|
|
|
|
queens_file = args.season_dir / QUEENS_FILE
|
|
picks_files = list(args.season_dir.glob(PICKS_FILE_GLOB))
|
|
eliminations_file = args.season_dir / ELIMINATIONS_FILE
|
|
rosters_file = args.season_dir / ROSTERS_FILE
|
|
|
|
if not eliminations_file.exists():
|
|
eliminations_file.touch()
|
|
|
|
draft = Draft(queens_file, picks_files, rosters_file)
|
|
|
|
if args.draft:
|
|
draft.run()
|
|
draft.save(rosters_file)
|
|
return
|
|
|
|
if args.eliminate:
|
|
eliminations = eliminate_queen(queens_file, eliminations_file, args.eliminate)
|
|
if eliminations:
|
|
draft.score_rosters(eliminations)
|
|
|
|
if args.print:
|
|
eliminations = eliminations_file.read_text().strip().split("\n")
|
|
print("Queens:")
|
|
print("\n".join(draft.queens))
|
|
print("Eliminations:")
|
|
print("\n".join(eliminations))
|
|
print("Rosters:")
|
|
draft.print_rosters()
|
|
print("Scores:")
|
|
draft.score_rosters(eliminations)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|