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()