Allow filtering Mega, Shadow, and Legendary Pokémon

src/pogo_scaled_estimators/calculator.py

@@ -17,6 +17,9 @@ from pogo_scaled_estimators.utilities import format_move_name, format_pokemon_na
 class Filter(Flag):
     NO_FILTER = 0
     DISALLOW_LEGACY_MOVES = auto()
+    DISALLOW_MEGA_POKEMON = auto()
+    DISALLOW_SHADOW_POKEMON = auto()
+    DISALLOW_LEGENDARY_POKEMON = auto()
 
 
 @final
@@ -32,6 +35,7 @@ class Calculator:
         attackers = {
             attacker: {"RAID_LEVEL_3": [], "RAID_LEVEL_5": [], "RAID_LEVEL_MEGA": []}
             for attacker in self._pokebattler_proxy.with_charged_moves(self.attacker_types)
+            if self._allowed_attacker(attacker)
         }
 
         with self._progress:
@@ -51,7 +55,9 @@ class Calculator:
                     self._progress.update(task, advance=1)
                     raid = Raid(raid_tier, defender, level, party)
                     results = {
-                        attacker: [moveset for moveset in movesets if self._viable_move(attacker, moveset.charged_move)]
+                        attacker: [
+                            moveset for moveset in movesets if self._allowed_move(attacker, moveset.charged_move)
+                        ]
                         for attacker, movesets in self._pokebattler_proxy.simulate(raid).items()
                         if attacker in attackers
                     }
@@ -82,7 +88,21 @@ class Calculator:
                 f"[bold]{format_pokemon_name(attacker, attacker_type)}[/bold] ({format_move_name(fast_move, fast_move_type)}/{format_move_name(charged_move, charged_move_type)}): {ase:.2f}"
             )
 
-    def _viable_move(self, pokemon_id: str, move_id: str) -> bool:
+    def _allowed_attacker(self, pokemon_id: str) -> bool:
+        if Filter.DISALLOW_MEGA_POKEMON in self.filters and ("_MEGA" in pokemon_id or "_PRIMAL" in pokemon_id):
+            return False
+        if Filter.DISALLOW_SHADOW_POKEMON in self.filters and "_SHADOW_FORM" in pokemon_id:
+            return False
+        pokemon = self._pokebattler_proxy.find_pokemon(pokemon_id)
+        if Filter.DISALLOW_LEGENDARY_POKEMON in self.filters and pokemon.get("rarity") in [
+            "POKEMON_RARITY_LEGENDARY",
+            "POKEMON_RARITY_MYTHIC",
+            "POKEMON_RARITY_ULTRA_BEAST",
+        ]:
+            return False
+        return True
+
+    def _allowed_move(self, pokemon_id: str, move_id: str) -> bool:
         move = self._pokebattler_proxy.find_move(move_id)
         if move.typing not in self.attacker_types:
             return False

src/pogo_scaled_estimators/cli.py

@@ -5,7 +5,9 @@
 # https://opensource.org/licenses/MIT.
 
 import argparse
+import operator
 import sys
+from functools import reduce
 
 from pogo_scaled_estimators.calculator import Calculator, Filter
 
@@ -15,12 +17,15 @@ def main_cli():
     _ = parser.add_argument("type", nargs="+", help="an attacker type")
     _ = parser.add_argument("--level", type=int, default=40)
     _ = parser.add_argument("--party", type=int, default=1)
-    _ = parser.add_argument("--no-elite", action="store_true")
+    _ = parser.add_argument("--no-legacy", dest="filters", action="append_const", const=Filter.DISALLOW_LEGACY_MOVES)
+    _ = parser.add_argument("--no-mega", dest="filters", action="append_const", const=Filter.DISALLOW_MEGA_POKEMON)
+    _ = parser.add_argument("--no-shadow", dest="filters", action="append_const", const=Filter.DISALLOW_SHADOW_POKEMON)
+    _ = parser.add_argument(
+        "--no-legendary", dest="filters", action="append_const", const=Filter.DISALLOW_LEGENDARY_POKEMON
+    )
 
     args = parser.parse_args()
-    filters = Filter.NO_FILTER
-    if args.no_elite:
-        filters |= Filter.DISALLOW_LEGACY_MOVES
+    filters = reduce(operator.or_, args.filters, Filter.NO_FILTER)
     calculator = Calculator(args.type, filters)
     calculator.calculate(level=args.level, party=args.party)