got strict tests to pass once

Author: cameronangliss

src/poke_env/battle/abstract_battle.py

@@ -319,7 +319,9 @@ def _check_damage_message_for_ability(self, split_message: List[str]):
         ):
             ability = split_message[4].split("ability:")[-1]
             pkmn = split_message[5].split("[of]")[-1].strip()
-            self.get_pokemon(pkmn).ability = to_id_str(ability)
+            pkmn_object = self.get_pokemon(pkmn)
+            if pkmn_object.ability is None:
+                pkmn_object.ability = ability
 
     def _check_heal_message_for_item(self, split_message: List[str]):
         # Catches when a side heals from it's own item
@@ -346,10 +348,14 @@ def _check_heal_message_for_ability(self, split_message: List[str]):
             ability = to_id_str(split_message[4].split("ability:")[-1])
             if ability == "hospitality":
                 pkmn = split_message[5].replace("[of] ", "").strip()
-                self.get_pokemon(pkmn).ability = ability
+                pkmn_object = self.get_pokemon(pkmn)
+                if pkmn_object.ability is None:
+                    pkmn_object.ability = ability
             else:
                 pkmn = split_message[2]
-                self.get_pokemon(pkmn).ability = ability
+                pkmn_object = self.get_pokemon(pkmn)
+                if pkmn_object.ability is None:
+                    pkmn_object.ability = ability
 
     @abstractmethod
     def end_illusion(self, pokemon_name: str, details: str):
@@ -529,7 +535,9 @@ def parse_message(self, split_message: List[str]):
                 revealed_ability = event.pop().split(": ")[-1]
 
                 pokemon = event[2]
-                self.get_pokemon(pokemon).ability = revealed_ability
+                mon = self.get_pokemon(pokemon)
+                if mon.ability is None:
+                    mon.ability = revealed_ability
 
                 if revealed_ability == "Magic Bounce":
                     return
@@ -672,6 +680,13 @@ def parse_message(self, split_message: List[str]):
             self.get_pokemon(pokemon).boost(stat, -int(amount))
         elif event[1] == "-ability":
             pokemon, cause = event[2:4]
+            mon = self.get_pokemon(pokemon)
+            if (
+                mon.ability is not None
+                and mon.ability in ["asoneglastrier", "asonespectrier"]
+                and event[3] in ["asone", "unnerve"]
+            ):
+                return
             if (
                 len(event) > 4
                 and (
@@ -685,11 +700,13 @@ def parse_message(self, split_message: List[str]):
                     or event[5].startswith("[from] ability: Trace")
                 )
             ):
-                self.get_pokemon(pokemon).set_temporary_ability(cause)
+                if mon.ability is None or mon._ability == to_id_str(cause):
+                    mon.ability = "trace"
+                mon.temporary_ability = cause
             elif cause == "Neutralizing Gas":
                 self.field_start(cause)
-            else:
-                self.get_pokemon(pokemon).ability = cause
+            elif mon.ability is None:
+                mon.ability = cause
         elif split_message[1] == "-start":
             pokemon, effect = event[2:4]
             mon = self.get_pokemon(pokemon)
@@ -720,10 +737,15 @@ def parse_message(self, split_message: List[str]):
             if target and effect == "move: Skill Swap":
                 self.get_pokemon(target).start_effect(effect, event[4:6])
                 actor = event[6].replace("[of] ", "")
-                self.get_pokemon(actor).set_temporary_ability(event[5])
+                self.get_pokemon(actor).temporary_ability = event[5]
             elif effect == "ability: Mummy":
                 target = event[5].replace("[of] ", "")
-                self.get_pokemon(target).set_temporary_ability("mummy")
+                self.get_pokemon(target).temporary_ability = "mummy"
+            elif effect == "ability: Wandering Spirit":
+                actor = event[2]
+                target = event[6].replace("[of] ", "")
+                self.get_pokemon(actor).temporary_ability = event[4]
+                self.get_pokemon(target).temporary_ability = "wanderingspirit"
             elif effect == "ability: Symbiosis":
                 self.get_pokemon(event[5].replace("[of] ", "")).item = event[4].replace(
                     "[item] ", ""
@@ -769,7 +791,7 @@ def parse_message(self, split_message: List[str]):
                 self.get_pokemon(pokemon).end_effect(effect)
         elif event[1] == "-endability":
             pokemon = event[2]
-            self.get_pokemon(pokemon).set_temporary_ability(None)
+            self.get_pokemon(pokemon).temporary_ability = None
         elif event[1] == "-enditem":
             pokemon, item = event[2:4]
             self.get_pokemon(pokemon).end_item(item)
@@ -801,22 +823,27 @@ def parse_message(self, split_message: List[str]):
                         elif mon == self.opponent_active_pokemon:
                             self.active_pokemon.item = to_id_str(item)
 
-                    mon.ability = to_id_str("frisk")
+                    if mon.ability is None:
+                        mon.ability = "frisk"
                 elif cause == "[from] ability: Pickpocket":
                     pickpocket = event[2]
                     pickpocketed = event[5].replace("[of] ", "")
                     item = event[3]
 
-                    self.get_pokemon(pickpocket).item = to_id_str(item)
-                    self.get_pokemon(pickpocket).ability = to_id_str("pickpocket")
+                    pickpocket = self.get_pokemon(pickpocket)
+                    pickpocket.item = to_id_str(item)
+                    if pickpocket.ability is None:
+                        pickpocket.ability = "pickpocket"
                     self.get_pokemon(pickpocketed).item = None
                 elif cause == "[from] ability: Magician":
                     magician = event[2]
                     victim = event[5].replace("[of] ", "")
                     item = event[3]
 
-                    self.get_pokemon(magician).item = to_id_str(item)
-                    self.get_pokemon(magician).ability = to_id_str("magician")
+                    magician = self.get_pokemon(magician)
+                    magician.item = to_id_str(item)
+                    if magician.ability is None:
+                        magician.ability = "magician"
                     self.get_pokemon(victim).item = None
                 elif cause in {"[from] move: Thief", "[from] move: Covet"}:
                     thief = event[2]
@@ -994,7 +1021,9 @@ def parse_message(self, split_message: List[str]):
 
                 if cause.startswith("[from] ability:"):
                     cause = cause.replace("[from] ability:", "")
-                    self.get_pokemon(pokemon).ability = to_id_str(cause)
+                    mon = self.get_pokemon(pokemon)
+                    if mon.ability is None:
+                        mon.ability = cause
         elif event[1] == "-swapsideconditions":
             self._side_conditions, self._opponent_side_conditions = (
                 self._opponent_side_conditions,

src/poke_env/battle/pokemon.py

@@ -13,8 +13,6 @@
 from poke_env.stats import compute_raw_stats
 from poke_env.teambuilder.teambuilder_pokemon import TeambuilderPokemon
 
-_NO_ABILITY = "NO_ABILITY"
-
 
 class Pokemon:
     __slots__ = (
@@ -34,6 +32,7 @@ class Pokemon:
         "_last_request",
         "_level",
         "_max_hp",
+        "_mega_ability",
         "_moves",
         "_must_recharge",
         "_name",
@@ -123,6 +122,7 @@ def __init__(
         self._status: Optional[Status] = None
         self._status_counter: int = 0
         self._temporary_ability: Optional[str] = None
+        self._mega_ability: Optional[str] = None
         self._temporary_types: List[PokemonType] = []
 
         if request_pokemon:
@@ -269,6 +269,7 @@ def end_turn(self):
     def faint(self):
         self._current_hp = 0
         self._status = Status.FNT
+        self.temporary_ability = None
         self._clear_effects()
 
     def forme_change(self, species: str):
@@ -290,6 +291,7 @@ def mega_evolve(self, stone: str):
             if not species_id_str.endswith("mega")
             else species_id_str
         )
+        self.temporary_ability = None
         if mega_species in self._data.pokedex:
             self._update_from_pokedex(mega_species, store_species=False)
         elif stone[-1] in "XYxy":
@@ -398,12 +400,6 @@ def set_hp_status(self, hp_status: str, store=False):
         if store:
             self._stats["hp"] = self._max_hp
 
-    def set_temporary_ability(self, ability: Optional[str]):
-        if ability is not None:
-            self._temporary_ability = to_id_str(ability)
-        else:
-            self._temporary_ability = _NO_ABILITY
-
     def start_effect(self, effect_str: str, details: Optional[Any] = None):
         effect = Effect.from_showdown_message(effect_str)
         if effect not in self._effects:
@@ -419,11 +415,10 @@ def start_effect(self, effect_str: str, details: Optional[Any] = None):
             for type_ in details.split("/"):
                 self._temporary_types.append(PokemonType.from_name(type_))
         elif effect == Effect.SKILL_SWAP and details:
-            self.set_temporary_ability(details[0])
-
             # Skill Swap reveals a mon's ability
             if self.ability is None:
-                self._ability = to_id_str(details[1])
+                self.ability = details[1]
+            self.temporary_ability = details[0]
 
     def _swap_boosts(self):
         self._boosts["atk"], self._boosts["spa"] = (
@@ -459,7 +454,7 @@ def switch_out(self, fields: Dict[Field, int]):
         self._preparing_move = None
         self._preparing_target = None
         self._protect_counter = 0
-        self._temporary_ability = None
+        self.temporary_ability = None
         self._temporary_types = []
 
         if self._status == Status.TOX:
@@ -489,12 +484,16 @@ def _update_from_pokedex(self, species: str, store_species: bool = True):
         else:
             self._type_2 = PokemonType.from_name(dex_entry["types"][1])
 
-        self._possible_abilities = [
-            to_id_str(ability) for ability in dex_entry["abilities"].values()
-        ]
-
-        if len(self._possible_abilities) == 1:
-            self.ability = self._possible_abilities[0]
+        if "forme" in dex_entry and dex_entry["forme"].startswith("Mega"):
+            self.mega_ability = dex_entry["abilities"]["0"]
+        elif self.mega_ability is None:
+            self._possible_abilities = [
+                to_id_str(ability) for ability in dex_entry["abilities"].values()
+            ]
+            if len(self._possible_abilities) == 1:
+                self.ability = self._possible_abilities[0]
+        else:
+            self.mega_ability = None
 
         self._heightm = dex_entry["heightm"]
         self._weightkg = dex_entry["weightkg"]
@@ -552,10 +551,13 @@ def update_from_request(self, request_pokemon: Dict[str, Any]):
         if request_pokemon == self._last_request:
             return
 
-        if "ability" in request_pokemon:
-            self.ability = request_pokemon["ability"]
-        elif "baseAbility" in request_pokemon:
+        if self.ability is None:
             self.ability = request_pokemon["baseAbility"]
+        if (
+            "ability" in request_pokemon
+            and request_pokemon["ability"] != request_pokemon["baseAbility"]
+        ):
+            self.temporary_ability = request_pokemon["ability"]
 
         self._last_request = request_pokemon
 
@@ -634,6 +636,15 @@ def check_consistency(self, pkmn_request: Dict[str, Any], player_role: str):
             assert Move.retrieve_id(move_request) == Move.retrieve_id(
                 move.id
             ), f"{Move.retrieve_id(move_request)} != {Move.retrieve_id(move.id)}\nrequest: {pkmn_request}"
+        if self.ability is None:
+            self.ability = pkmn_request["baseAbility"]
+        assert pkmn_request["baseAbility"] == (
+            self.base_ability or ""
+        ), f"{pkmn_request['baseAbility']} != {self.base_ability or ''}"
+        if "ability" in pkmn_request:
+            assert pkmn_request["ability"] == (
+                self.ability or ""
+            ), f"{pkmn_request['ability']} != {self.ability or ''}"
 
     def _update_from_teambuilder(self, tb: TeambuilderPokemon):
         if tb.nickname is not None and tb.species is None:
@@ -760,19 +771,16 @@ def ability(self) -> Optional[str]:
         :return: The pokemon's ability. None if unknown or removed.
         :rtype: str, optional
         """
-        if self._temporary_ability == _NO_ABILITY:
-            return None
-        elif self._temporary_ability is not None:
+        if self._temporary_ability is not None:
             return self._temporary_ability
+        elif self.mega_ability is not None:
+            return self.mega_ability
         else:
             return self._ability
 
     @ability.setter
-    def ability(self, ability: Optional[str]):
-        if ability is None:
-            self._ability = None
-        else:
-            self._ability = to_id_str(ability)
+    def ability(self, ability: str):
+        self._ability = to_id_str(ability)
 
     @property
     def active(self) -> Optional[bool]:
@@ -806,6 +814,14 @@ def available_z_moves(self) -> List[Move]:
                 return [self._moves[move]]
         return []
 
+    @property
+    def base_ability(self) -> Optional[str]:
+        """
+        :return: The pokemon's base ability. None if unknown.
+        :rtype: str, optional
+        """
+        return self._mega_ability or self._ability
+
     @property
     def base_species(self) -> str:
         """
@@ -961,6 +977,18 @@ def max_hp(self) -> int:
         """
         return self._max_hp or 0
 
+    @property
+    def mega_ability(self) -> Optional[str]:
+        """
+        :return: The pokemon's mega ability. None if the pokemon is not mega evolved.
+        :rtype: str, optional
+        """
+        return self._mega_ability
+
+    @mega_ability.setter
+    def mega_ability(self, ability: Optional[str]):
+        self._mega_ability = to_id_str(ability) if ability is not None else None
+
     @property
     def moves(self) -> Dict[str, Move]:
         """
@@ -1135,6 +1163,18 @@ def stab_multiplier(self) -> float:
             return 2
         return 1.5
 
+    @property
+    def temporary_ability(self) -> str | None:
+        """
+        :return: The pokemon's temporary ability, None if none is set.
+        :rtype: Optional[str]
+        """
+        return self._temporary_ability
+
+    @temporary_ability.setter
+    def temporary_ability(self, ability: str | None):
+        self._temporary_ability = to_id_str(ability) if ability is not None else None
+
     @property
     def tera_type(self) -> Optional[PokemonType]:
         """