Track regenerator healing correctly (#805)

* track regenerator and trace abilities properly

* fix

* track mummy correctly too

* fix test

* isort

* fix test

* fix test

* be careful around edge case

* fix test

* better solution

Author: cameronangliss

src/poke_env/battle/abstract_battle.py

@@ -378,14 +378,17 @@ def _end_illusion_on(
         )
         illusionist_mon.set_hp(f"{illusioned.current_hp}/{illusioned.max_hp}")
 
-        illusioned.was_illusioned()
+        illusioned.was_illusioned(self.fields)
 
         return illusionist_mon
 
     def _field_end(self, field_str: str):
         field = Field.from_showdown_message(field_str)
         if field is not Field.UNKNOWN:
-            self._fields.pop(field)
+            if field is Field.NEUTRALIZING_GAS:
+                self._fields.pop(field, 0)
+            else:
+                self._fields.pop(field)
 
     def field_start(self, field_str: str):
         field = Field.from_showdown_message(field_str)
@@ -676,8 +679,22 @@ 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]
-            if len(event) > 4 and event[4].startswith("[from] move:"):
+            if (
+                len(event) > 4
+                and (
+                    event[4].startswith("[from] move:")
+                    or event[4].startswith("[from] ability: Trace")
+                )
+            ) or (
+                len(event) > 5
+                and (
+                    event[5].startswith("[from] move:")
+                    or event[5].startswith("[from] ability: Trace")
+                )
+            ):
                 self.get_pokemon(pokemon).set_temporary_ability(cause)
+            elif cause == "Neutralizing Gas":
+                self.field_start(cause)
             else:
                 self.get_pokemon(pokemon).ability = cause
         elif split_message[1] == "-start":
@@ -711,6 +728,9 @@ def parse_message(self, split_message: List[str]):
                 self.get_pokemon(target).start_effect(effect, event[4:6])
                 actor = event[6].replace("[of] ", "")
                 self.get_pokemon(actor).set_temporary_ability(event[5])
+            elif effect == "ability: Mummy":
+                target = event[5].replace("[of] ", "")
+                self.get_pokemon(target).set_temporary_ability("mummy")
             elif effect == "ability: Symbiosis":
                 self.get_pokemon(event[5].replace("[of] ", "")).item = event[4].replace(
                     "[item] ", ""
@@ -750,7 +770,10 @@ def parse_message(self, split_message: List[str]):
                 mon.cure_status()
         elif event[1] == "-end":
             pokemon, effect = event[2:4]
-            self.get_pokemon(pokemon).end_effect(effect)
+            if effect == "ability: Neutralizing Gas":
+                self._field_end(effect)
+            else:
+                self.get_pokemon(pokemon).end_effect(effect)
         elif event[1] == "-endability":
             pokemon = event[2]
             self.get_pokemon(pokemon).set_temporary_ability(None)

src/poke_env/battle/battle.py

@@ -130,10 +130,10 @@ def switch(self, pokemon_str: str, details: str, hp_status: str):
 
         if identifier == self._player_role:
             if self.active_pokemon:
-                self.active_pokemon.switch_out()
+                self.active_pokemon.switch_out(self.fields)
         else:
             if self.opponent_active_pokemon:
-                self.opponent_active_pokemon.switch_out()
+                self.opponent_active_pokemon.switch_out(self.fields)
 
         pokemon = self.get_pokemon(pokemon_str, details=details)
 

src/poke_env/battle/double_battle.py

@@ -220,7 +220,7 @@ def switch(self, pokemon_str: str, details: str, hp_status: str):
         )
         pokemon_out = team.pop(pokemon_identifier, None)
         if pokemon_out is not None:
-            pokemon_out.switch_out()
+            pokemon_out.switch_out(self.fields)
         pokemon_in = self.get_pokemon(pokemon_str, details=details)
         pokemon_in.switch_in()
         pokemon_in.set_hp_status(hp_status)

src/poke_env/battle/effect.py

@@ -126,7 +126,6 @@ class Effect(Enum):
     MISTY_TERRAIN = auto()
     MUMMY = auto()
     MUST_RECHARGE = auto()
-    NEUTRALIZING_GAS = auto()
     NIGHTMARE = auto()
     NO_RETREAT = auto()
     OBLIVIOUS = auto()
@@ -406,7 +405,6 @@ def is_from_move(self) -> bool:
     Effect.LIQUID_OOZE,
     Effect.MIMICRY,
     Effect.MUMMY,
-    Effect.NEUTRALIZING_GAS,
     Effect.OBLIVIOUS,
     Effect.ORICHALCUM_PULSE,
     Effect.OWN_TEMPO,
@@ -891,7 +889,6 @@ def is_from_move(self) -> bool:
     "MISTYTERRAIN": Effect.MISTY_TERRAIN,
     "MUMMY": Effect.MUMMY,
     "MUSTRECHARGE": Effect.MUST_RECHARGE,
-    "NEUTRALIZINGGAS": Effect.NEUTRALIZING_GAS,
     "NIGHTMARE": Effect.NIGHTMARE,
     "NORETREAT": Effect.NO_RETREAT,
     "OBLIVIOUS": Effect.OBLIVIOUS,

src/poke_env/battle/field.py

@@ -20,6 +20,7 @@ class Field(Enum):
     MISTY_TERRAIN = auto()
     MUD_SPORT = auto()
     MUD_SPOT = auto()
+    NEUTRALIZING_GAS = auto()
     PSYCHIC_TERRAIN = auto()
     TRICK_ROOM = auto()
     WATER_SPORT = auto()
@@ -37,6 +38,7 @@ def from_showdown_message(message: str) -> Field:
         :return: The corresponding Field object.
         :rtype: Field
         """
+        message = message.replace("ability: ", "")
         message = message.replace("move: ", "")
         message = message.replace(" ", "_")
 

src/poke_env/battle/pokemon.py

@@ -3,6 +3,7 @@
 from typing import Any, Dict, List, Optional, Union
 
 from poke_env.battle.effect import Effect
+from poke_env.battle.field import Field
 from poke_env.battle.move import SPECIAL_MOVES, Move
 from poke_env.battle.pokemon_gender import PokemonGender
 from poke_env.battle.pokemon_type import PokemonType
@@ -438,7 +439,17 @@ def switch_in(self, details: Optional[str] = None):
         self._first_turn = True
         self._revealed = True
 
-    def switch_out(self):
+    def switch_out(self, fields: Dict[Field, int]):
+        if (
+            self.ability == "regenerator"
+            and (
+                self.item == "abilityshield"
+                or Field.NEUTRALIZING_GAS not in fields.keys()
+            )
+            and self.status != Status.FNT
+        ):
+            self._current_hp = min(int(self.current_hp + self.max_hp / 3), self.max_hp)
+
         self._active = False
         self.clear_boosts()
         self._clear_effects()
@@ -447,17 +458,12 @@ def switch_out(self):
         self._preparing_move = None
         self._preparing_target = None
         self._protect_counter = 0
+        self._temporary_ability = None
+        self._temporary_types = []
 
         if self._status == Status.TOX:
             self._status_counter = 0
 
-        for effect in self.effects:
-            if effect.ends_on_switch or effect.is_volatile_status:
-                self.end_effect(effect.name)
-
-        self._temporary_ability = None
-        self._temporary_types = []
-
     def terastallize(self, type_: str):
         self._terastallized_type = PokemonType.from_name(type_)
         self._terastallized = True
@@ -667,7 +673,7 @@ def _update_from_teambuilder(self, tb: TeambuilderPokemon):
     def used_z_move(self):
         self._item = None
 
-    def was_illusioned(self):
+    def was_illusioned(self, fields: Dict[Field, int]):
         self._current_hp = None
         self._max_hp = None
         self._status = None
@@ -678,7 +684,7 @@ def was_illusioned(self):
         if last_request:
             self.update_from_request(last_request)
 
-        self.switch_out()
+        self.switch_out(fields)
 
     def available_moves_from_request(self, request: Dict[str, Any]) -> List[Move]:
         moves: List[Move] = []

unit_tests/environment/test_battle.py

@@ -545,7 +545,7 @@ def test_battle_request_and_interactions(example_request):
     necrozma = battle.active_pokemon
     groudon = battle.opponent_active_pokemon
 
-    necrozma.switch_out()
+    necrozma.switch_out(battle.fields)
     groudon.switch_in()
     groudon._ability = None
 
@@ -595,7 +595,7 @@ def test_battle_request_and_interactions(example_request):
         ["", "-ability", "p1a: Groudon", "Insomnia", "[from] move: Worry Seed"]
     )
     assert groudon.ability == "insomnia"
-    groudon.switch_out()
+    groudon.switch_out(battle.fields)
     groudon.switch_in()
     assert groudon.ability == "desolateland"
 
@@ -614,7 +614,7 @@ def test_battle_request_and_interactions(example_request):
     assert groudon.ability == "prismarmor"
     assert necrozma.ability == "desolateland"
     groudon.switch_in()
-    groudon.switch_out()
+    groudon.switch_out(battle.fields)
     assert groudon.ability == "desolateland"
 
     battle.parse_message(["", "switch", "p1a: Ho-oh", "Ho-oh, L82", "100/100"])
@@ -631,7 +631,7 @@ def test_battle_request_and_interactions(example_request):
     assert hooh.type_1 == PokemonType.THREE_QUESTION_MARKS
     assert hooh.type_2 == PokemonType.FLYING
     assert hooh.types == [PokemonType.THREE_QUESTION_MARKS, PokemonType.FLYING]
-    hooh.switch_out()
+    hooh.switch_out(battle.fields)
     hooh.switch_in()
     assert hooh.type_1 == PokemonType.FIRE
     assert hooh.type_2 == PokemonType.FLYING
@@ -642,7 +642,7 @@ def test_battle_request_and_interactions(example_request):
     assert hooh.type_1 == PokemonType.WATER
     assert hooh.type_2 is None
     assert hooh.types == [PokemonType.WATER]
-    hooh.switch_out()
+    hooh.switch_out(battle.fields)
     hooh.switch_in()
     assert hooh.type_1 == PokemonType.FIRE
     assert hooh.type_2 == PokemonType.FLYING

unit_tests/environment/test_enumerations.py

@@ -115,7 +115,7 @@ def test_effect_end():
     assert Effect.QUASH not in furret.effects
 
     # Test end on switch
-    furret.switch_out()
+    furret.switch_out({})
     assert Effect.MUMMY not in furret.effects
 
     # Test the definition of Volatile Status

unit_tests/environment/test_pokemon.py

@@ -332,7 +332,7 @@ def test_temporary():
     furret.start_effect("typechange", "Fighting")
     assert furret.damage_multiplier(PokemonType.PSYCHIC) == 2
 
-    furret.switch_out()
+    furret.switch_out({})
     furret.switch_in()
 
     assert furret.ability == "adaptability"
@@ -348,6 +348,6 @@ def test_temporary():
     furret.set_temporary_ability(None)
     assert furret.ability is None
 
-    furret.switch_out()
+    furret.switch_out({})
     assert furret.ability == "adaptability"
     assert furret.types == [PokemonType.DRAGON]