Coverage for src/duelboard/utils.py: 100%

1"""Utility functions for Elo calculations.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING 

6 

7import pandas as pd 

8 

9from .models import Battle, BattleOutcome 

10 

11if TYPE_CHECKING: 

12 from .types import RatingsDict 

13 

14 

15def load_battles_from_csv( 

16 file_path: str, 

17 player_a_col: str = "player_a", 

18 player_b_col: str = "player_b", 

19 winner_col: str = "winner", 

20 **pandas_kwargs: object, 

21) -> list[Battle]: 

22 """Load battles from CSV file. 

23 

24 Args: 

25 file_path: Path to CSV file 

26 player_a_col: Column name for player A 

27 player_b_col: Column name for player B 

28 winner_col: Column name for winner 

29 **pandas_kwargs: Additional arguments for pandas.read_csv 

30 

31 Returns: 

32 List of Battle objects 

33 """ 

34 df = pd.read_csv(file_path, **pandas_kwargs) 

35 return dataframe_to_battles(df, player_a_col, player_b_col, winner_col) 

36 

37 

38def load_battles_from_json( 

39 file_path: str, 

40 player_a_col: str = "player_a", 

41 player_b_col: str = "player_b", 

42 winner_col: str = "winner", 

43 **pandas_kwargs: object, 

44) -> list[Battle]: 

45 """Load battles from JSON file. 

46 

47 Args: 

48 file_path: Path to JSON file 

49 player_a_col: Column name for player A 

50 player_b_col: Column name for player B 

51 winner_col: Column name for winner 

52 **pandas_kwargs: Additional arguments for pandas.read_json 

53 

54 Returns: 

55 List of Battle objects 

56 """ 

57 df = pd.read_json(file_path, **pandas_kwargs) 

58 return dataframe_to_battles(df, player_a_col, player_b_col, winner_col) 

59 

60 

61def dataframe_to_battles( 

62 df: pd.DataFrame, 

63 player_a_col: str = "player_a", 

64 player_b_col: str = "player_b", 

65 winner_col: str = "winner", 

66) -> list[Battle]: 

67 """Convert DataFrame to list of Battle objects. 

68 

69 Args: 

70 df: DataFrame with battle data 

71 player_a_col: Column name for player A 

72 player_b_col: Column name for player B 

73 winner_col: Column name for winner 

74 

75 Returns: 

76 List of Battle objects 

77 """ 

78 battles = [] 

79 for _, row in df.iterrows(): 

80 # Create metadata from remaining columns 

81 metadata = row.drop([player_a_col, player_b_col, winner_col]).to_dict() 

82 

83 battle = Battle( 

84 player_a=row[player_a_col], 

85 player_b=row[player_b_col], 

86 outcome=BattleOutcome(row[winner_col]), 

87 metadata=metadata, 

88 ) 

89 battles.append(battle) 

90 

91 return battles 

92 

93 

94def battles_to_dataframe(battles: list[Battle]) -> pd.DataFrame: 

95 """Convert list of Battle objects to DataFrame. 

96 

97 Args: 

98 battles: List of Battle objects 

99 

100 Returns: 

101 DataFrame with battle data 

102 """ 

103 data = [] 

104 for battle in battles: 

105 row = { 

106 "player_a": battle.player_a, 

107 "player_b": battle.player_b, 

108 "winner": battle.outcome.value, 

109 } 

110 

111 # Add metadata if available 

112 if battle.metadata: 

113 row.update(battle.metadata) 

114 

115 data.append(row) 

116 

117 return pd.DataFrame(data) 

118 

119 

120def filter_anonymous_battles(df: pd.DataFrame, anony_col: str = "anony") -> pd.DataFrame: 

121 """Filter to only anonymous battles (following the notebook approach). 

122 

123 Args: 

124 df: DataFrame with battle data 

125 anony_col: Column name for anonymity flag 

126 

127 Returns: 

128 Filtered DataFrame with only anonymous battles 

129 """ 

130 if anony_col in df.columns: 

131 return df[df[anony_col]].reset_index(drop=True) 

132 return df 

133 

134 

135def filter_non_tie_battles(df: pd.DataFrame, winner_col: str = "winner") -> pd.DataFrame: 

136 """Filter out tie battles. 

137 

138 Args: 

139 df: DataFrame with battle data 

140 winner_col: Column name for winner 

141 

142 Returns: 

143 Filtered DataFrame without tie battles 

144 """ 

145 return df[~df[winner_col].str.contains("tie", case=False, na=False)].reset_index(drop=True) 

146 

147 

148def get_rating_summary(ratings: RatingsDict) -> pd.DataFrame: 

149 """Get summary DataFrame of ratings. 

150 

151 Args: 

152 ratings: Dictionary of player ratings 

153 

154 Returns: 

155 DataFrame with rating summary sorted by rating 

156 """ 

157 data = [] 

158 for player, rating in ratings.items(): 

159 row = { 

160 "player": player, 

161 "rating": rating.rating, 

162 "battles": rating.battles, 

163 } 

164 

165 if rating.confidence_interval: 

166 row["ci_lower"] = rating.confidence_interval[0] 

167 row["ci_upper"] = rating.confidence_interval[1] 

168 row["ci_width"] = rating.confidence_interval[1] - rating.confidence_interval[0] 

169 

170 data.append(row) 

171 

172 df = pd.DataFrame(data) 

173 return df.sort_values("rating", ascending=False).reset_index(drop=True) 

174 

175 

176def rank_players_by_rating(ratings: RatingsDict) -> list[str]: 

177 """Get player names ranked by rating (highest first). 

178 

179 Args: 

180 ratings: Dictionary of player ratings 

181 

182 Returns: 

183 List of player names sorted by rating 

184 """ 

185 return [rating.player for rating in sorted(ratings.values(), key=lambda x: x.rating, reverse=True)] 

186 

187 

188def export_ratings_to_csv(ratings: RatingsDict, file_path: str) -> None: 

189 """Export ratings to CSV file. 

190 

191 Args: 

192 ratings: Dictionary of player ratings 

193 file_path: Path to output CSV file 

194 """ 

195 df = get_rating_summary(ratings) 

196 df.to_csv(file_path, index=False)