diff --git a/src/pogo_scaled_estimators/calculator.py b/src/pogo_scaled_estimators/calculator.py index 3377fc9..c40c9e4 100644 --- a/src/pogo_scaled_estimators/calculator.py +++ b/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 diff --git a/src/pogo_scaled_estimators/cli.py b/src/pogo_scaled_estimators/cli.py index 4d4bbad..fe8bf45 100644 --- a/src/pogo_scaled_estimators/cli.py +++ b/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)