Pretty colors everywhere

src/pogo_scaled_estimators/calculator.py

@@ -28,7 +28,7 @@ class Calculator:
 
         with self._progress:
             total_defenders = sum(len(defenders) for defenders in raid_bosses.values())
-            task = self._progress.add_task("Working...", total=(total_defenders * 30))
+            task = self._progress.add_task("Simulating raids...", total=total_defenders)
 
             for raid_tier, defenders in raid_bosses.items():
                 if raid_tier == "RAID_LEVEL_MEGA_5":
@@ -38,7 +38,9 @@ class Calculator:
                 else:
                     simplified_raid_tier = raid_tier
                 for defender in defenders:
-                    self._progress.update(task, description=f"vs {format_pokemon_name(defender)}...")
+                    defender_type = self._pokebattler_proxy.pokemon_type(defender)
+                    self._progress.console.log(f"vs {format_pokemon_name(defender, defender_type)}...")
+                    self._progress.update(task, advance=1)
                     raid = Raid(raid_tier, defender, level, party)
                     results = {
                         attacker: movesets
@@ -51,10 +53,9 @@ class Calculator:
                     }
                     best_estimator = min(best_movesets.values(), key=lambda moveset: moveset.estimator).estimator
                     for attacker, moveset in best_movesets.items():
-                        self._progress.update(task, advance=1)
                         attackers[attacker][simplified_raid_tier].append(moveset.scale(best_estimator))
 
-        ase = {}
+        ase_by_attacker: list[tuple[str, float]] = []
         for attacker, raid_tier_results in attackers.items():
             if (
                 not raid_tier_results["RAID_LEVEL_3"]
@@ -67,7 +68,11 @@ class Calculator:
                 + 0.50 * self._average_estimator(raid_tier_results["RAID_LEVEL_5"])
                 + 0.35 * self._average_estimator(raid_tier_results["RAID_LEVEL_MEGA"])
             )
-            print(f"{attacker},{level},{ase}")
+            ase_by_attacker.append((attacker, ase))
+        ase_by_attacker.sort(key=lambda item: item[1])
+        for attacker, ase in ase_by_attacker:
+            attacker_type = self._pokebattler_proxy.pokemon_type(attacker)
+            self._progress.console.print(f"{format_pokemon_name(attacker, attacker_type)}: {ase:.2f}")
 
     def _average_estimator(self, movesets: list[Moveset]) -> float:
         return sum(moveset.estimator for moveset in movesets) / len(movesets)

src/pogo_scaled_estimators/pokebattler_proxy.py

@@ -151,3 +151,6 @@ class PokebattlerProxy:
             for mon in self.pokemon
             if any(moveset["cinematicMove"] in charged_moves for moveset in mon["movesets"])
         ]
+
+    def pokemon_type(self, name: str) -> str | None:
+        return next(filter(lambda mon: mon["pokemonId"] == name, self.pokemon))["type"]

src/pogo_scaled_estimators/utilities.py

@@ -1,7 +1,28 @@
 MINIMUM_SPECIAL_NAME_PARTS = 2
 
 
-def format_pokemon_name(name):
+POKEMON_TYPE_COLORS = {
+    "POKEMON_TYPE_BUG": "green_yellow",
+    "POKEMON_TYPE_DARK": "bright_black",
+    "POKEMON_TYPE_DRAGON": "dodger_blue2",
+    "POKEMON_TYPE_ELECTRIC": "yellow1",
+    "POKEMON_TYPE_FAIRY": "orchid1",
+    "POKEMON_TYPE_FIGHTING": "red3",
+    "POKEMON_TYPE_FIRE": "orange1",
+    "POKEMON_TYPE_FLYING": "light_sky_blue1",
+    "POKEMON_TYPE_GHOST": "slate_blue3",
+    "POKEMON_TYPE_GRASS": "green3",
+    "POKEMON_TYPE_GROUND": "orange4",
+    "POKEMON_TYPE_ICE": "pale_turquoise1",
+    "POKEMON_TYPE_NORMAL": "grey53",
+    "POKEMON_TYPE_POISON": "dark_magenta",
+    "POKEMON_TYPE_PSYCHIC": "hot_pink",
+    "POKEMON_TYPE_ROCK": "gold3",
+    "POKEMON_TYPE_STEEL": "steel_blue",
+    "POKEMON_TYPE_WATER": "cornflower_blue",
+}
+
+def format_pokemon_name(name: str, pokemon_type: str | None = None):
     parts = [part.capitalize() for part in name.split("_")]
     if parts[-1] == "Mega" or parts[-1] == "Primal":
         parts = [parts[-1]] + parts[:-1]
@@ -9,7 +30,10 @@ def format_pokemon_name(name):
         parts = [parts[-2]] + parts[:-2] + [parts[-1]]
     if len(parts) > MINIMUM_SPECIAL_NAME_PARTS and parts[-2] == "Shadow":
         parts = [parts[-2]] + parts[:-2]
-    return " ".join(parts)
+    formatted_name = " ".join(parts)
+    if pokemon_type:
+        return f"[{POKEMON_TYPE_COLORS[pokemon_type]}]{formatted_name}[/]"
+    return formatted_name
 
 
 def format_move_name(name):