Repository: hsahovic/poke-env

Issue:
Add support for typechange and skillswap
Right now typechange and skillswap is a non-niche interaction in VGC that is not handleable in poke-env today. Planning to:
- modify the error of Effect.TYPE_CHANGE -> Effect.TYPECHANGE
- add support for typechange using a new temporary_types internal field
- add support for skill swap using a new temporary_ability field
- edit end_effect to also undo the temporary changes above
- change clear_effects to end effect each effect of the pokemon, and not to just reset the map
- change handling of protosynthesis/quark drive too since we store multiple effects due to duplicate showdown protocols. We should remove them all when one gets removed

Discussion:

Solution patch:
diff --git a/src/poke_env/environment/abstract_battle.py b/src/poke_env/environment/abstract_battle.py
index 8a342c82a..fa33333ed 100644
--- a/src/poke_env/environment/abstract_battle.py
+++ b/src/poke_env/environment/abstract_battle.py
@@ -610,11 +610,18 @@ def parse_message(self, split_message: List[str]):
             self.get_pokemon(pokemon).boost(stat, -int(amount))
         elif event[1] == "-ability":
             pokemon, ability = event[2:4]
-            self.get_pokemon(pokemon).ability = ability
-        elif event[1] == "-start":
+            if len(event) > 4 and event[4].startswith("[from] move:"):
+                self.get_pokemon(pokemon).set_temporary_ability(ability)
+            else:
+                self.get_pokemon(pokemon).ability = ability
+        elif split_message[1] == "-start":
             pokemon, effect = event[2:4]
-            pokemon = self.get_pokemon(pokemon)  # type: ignore
-            pokemon.start_effect(effect)  # type: ignore
+            pokemon = self.get_pokemon(pokemon)
+
+            if effect == "typechange":
+                pokemon.start_effect(effect, details=event[4])
+            else:
+                pokemon.start_effect(effect)
 
             if pokemon.is_dynamaxed:  # type: ignore
                 if pokemon in set(self.team.values()) and self._dynamax_turn is None:
@@ -628,7 +635,11 @@ def parse_message(self, split_message: List[str]):
                     self.opponent_can_dynamax = False
         elif event[1] == "-activate":
             target, effect = event[2:4]
-            if target:
+            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])
+            else:
                 self.get_pokemon(target).start_effect(effect)
         elif event[1] == "-status":
             pokemon, status = event[2:4]
@@ -665,7 +676,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).ability = None
+            self.get_pokemon(pokemon).set_temporary_ability(None)
         elif event[1] == "-enditem":
             pokemon, item = event[2:4]
             self.get_pokemon(pokemon).end_item(item)
diff --git a/src/poke_env/environment/pokemon.py b/src/poke_env/environment/pokemon.py
index 0cab4d304..35387fa36 100644
--- a/src/poke_env/environment/pokemon.py
+++ b/src/poke_env/environment/pokemon.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Union
 
 from poke_env.data import GenData, to_id_str
 from poke_env.environment.effect import Effect
@@ -12,6 +12,8 @@
 from poke_env.stats import compute_raw_stats
 from poke_env.teambuilder.teambuilder_pokemon import TeambuilderPokemon
 
+_NO_ABILITY = "NO_ABILITY"
+
 
 class Pokemon:
     __slots__ = (
@@ -44,6 +46,8 @@ class Pokemon:
         "_species",
         "_status",
         "_status_counter",
+        "_temporary_ability",
+        "_temporary_types",
         "_terastallized",
         "_terastallized_type",
         "_type_1",
@@ -117,6 +121,8 @@ def __init__(
         }
         self._status: Optional[Status] = None
         self._status_counter: int = 0
+        self._temporary_ability: Optional[str] = None
+        self._temporary_types: List[PokemonType] = []
 
         if request_pokemon:
             self.update_from_request(request_pokemon)
@@ -181,7 +187,8 @@ def clear_boosts(self):
             self._boosts[stat] = 0
 
     def _clear_effects(self):
-        self._effects = {}
+        for effect in [e for e in self._effects]:
+            self.end_effect(effect.name)
 
     def clear_negative_boosts(self):
         for stat, value in self._boosts.items():
@@ -211,6 +218,37 @@ def end_effect(self, effect_str: str):
         if effect in self._effects:
             self._effects.pop(effect)
 
+        if effect == Effect.TYPECHANGE:
+            self._temporary_types = []
+
+        elif effect == Effect.SKILL_SWAP:
+            self._temporary_ability = None
+
+        # Because of Showdown protocols, Protosynthesis and Quark Drive
+        # both get two Effects added to the mon, but only have one protocol
+        # message when the Effect ends
+        elif effect == Effect.QUARK_DRIVE:
+            for effect in [
+                Effect.QUARKDRIVEATK,
+                Effect.QUARKDRIVEDEF,
+                Effect.QUARKDRIVESPA,
+                Effect.QUARKDRIVESPD,
+                Effect.QUARKDRIVESPE,
+            ]:
+                if effect in self._effects:
+                    self._effects.pop(effect)
+
+        elif effect == Effect.PROTOSYNTHESIS:
+            for effect in [
+                Effect.PROTOSYNTHESISATK,
+                Effect.PROTOSYNTHESISDEF,
+                Effect.PROTOSYNTHESISSPA,
+                Effect.PROTOSYNTHESISSPD,
+                Effect.PROTOSYNTHESISSPE,
+            ]:
+                if effect in self._effects:
+                    self._effects.pop(effect)
+
     def end_item(self, item: str):
         self._item = None
 
@@ -230,7 +268,7 @@ def end_turn(self):
     def faint(self):
         self._current_hp = 0
         self._status = Status.FNT
-        self._effects = {}
+        self._clear_effects()
 
     def forme_change(self, species: str):
         species = species.split(",")[0]
@@ -346,7 +384,13 @@ def set_hp_status(self, hp_status: str, store=False):
         if store:
             self._stats["hp"] = self._max_hp
 
-    def start_effect(self, effect_str: str):
+    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:
             self._effects[effect] = 0
@@ -356,6 +400,17 @@ def start_effect(self, effect_str: str):
         if effect.breaks_protect:
             self._protect_counter = 0
 
+        if effect == Effect.TYPECHANGE and details:
+            self._temporary_types = []
+            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])
+
     def _swap_boosts(self):
         self._boosts["atk"], self._boosts["spa"] = (
             self._boosts["spa"],
@@ -388,9 +443,13 @@ def switch_out(self):
             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
+        self._temporary_types = []
 
     def transform(self, into: Pokemon):
         current_hp = self.current_hp
@@ -610,16 +669,21 @@ def damage_multiplier(self, type_or_move: Union[PokemonType, Move]) -> float:
         if isinstance(type_or_move, Move):
             type_or_move = type_or_move.type
         return type_or_move.damage_multiplier(
-            self._type_1, self._type_2, type_chart=self._data.type_chart
+            self.type_1, self.type_2, type_chart=self._data.type_chart
         )
 
     @property
     def ability(self) -> Optional[str]:
         """
-        :return: The pokemon's ability. None if unknown.
+        :return: The pokemon's ability. None if unknown or removed.
         :rtype: str, optional
         """
-        return self._ability
+        if self._temporary_ability == _NO_ABILITY:
+            return None
+        elif self._temporary_ability is not None:
+            return self._temporary_ability
+        else:
+            return self._ability
 
     @ability.setter
     def ability(self, ability: Optional[str]):
@@ -944,12 +1008,24 @@ def stab_multiplier(self) -> float:
         :return: The pokemon's STAB multiplier.
         :rtype: float
         """
-        if self._terastallized and self._terastallized_type in (
-            self._type_1,
-            self._type_2,
+        # https://bulbapedia.bulbagarden.net/wiki/Terastal_phenomenon#Effects
+        if (
+            self._terastallized
+            and (
+                self.tera_type in (self._type_1, self._type_2)
+                or self.tera_type in self._temporary_types
+            )
+            and self.ability == "adaptability"
         ):
+            return 2.25
+        elif self._terastallized and (
+            self.tera_type in (self._type_1, self._type_2)
+            or self.tera_type in self._temporary_types
+        ):
+            return 2
+        elif self.ability == "adaptability":
             return 2
-        return 1
+        return 1.5
 
     @property
     def tera_type(self) -> Optional[PokemonType]:
@@ -967,6 +1043,8 @@ def type_1(self) -> PokemonType:
         """
         if self._terastallized and self._terastallized_type is not None:
             return self._terastallized_type
+        elif len(self._temporary_types) > 0:
+            return self._temporary_types[0]
         return self._type_1
 
     @property
@@ -975,17 +1053,26 @@ def type_2(self) -> Optional[PokemonType]:
         :return: The pokemon's second type.
         :rtype: Optional[PokemonType]
         """
-        if self._terastallized:
+        if self._terastallized or len(self._temporary_types) == 1:
             return None
+        elif len(self._temporary_types) > 1:
+            return self._temporary_types[1]
         return self._type_2
 
     @property
-    def types(self) -> Tuple[PokemonType, Optional[PokemonType]]:
+    def types(self) -> List[PokemonType]:
         """
-        :return: The pokemon's types, as a tuple.
-        :rtype: Tuple[PokemonType, Optional[PokemonType]]
+        :return: The pokemon's types. Not a Tuple because moves can add types
+            and also remove types, and override them too.
+        :rtype: Tuple[PokemonType]
         """
-        return self.type_1, self._type_2
+        if len(self._temporary_types) > 0:
+            return self._temporary_types
+        else:
+            types = [self.type_1]
+            if self.type_2 is not None:
+                types.append(self.type_2)
+            return types
 
     @property
     def weight(self) -> float:


Test patch:
diff --git a/unit_tests/environment/test_battle.py b/unit_tests/environment/test_battle.py
index f7accdb3f..dafb74318 100644
--- a/unit_tests/environment/test_battle.py
+++ b/unit_tests/environment/test_battle.py
@@ -33,10 +33,10 @@ def test_battle_get_pokemon():
     assert "p1: hydreigon" in battle.opponent_team
 
     assert battle.get_pokemon("p2: tapufini").species == "tapufini"
-    assert battle.get_pokemon("p2: tapubulu").types == (
+    assert battle.get_pokemon("p2: tapubulu").types == [
         PokemonType.GRASS,
         PokemonType.FAIRY,
-    )
+    ]
     assert battle.get_pokemon("p2: tapulele").base_stats == {
         "atk": 85,
         "def": 75,
@@ -539,7 +539,12 @@ def test_battle_request_and_interactions(example_request):
         ]
     )
     assert battle.opponent_active_pokemon.ability == "waterabsorb"
-    battle.opponent_active_pokemon._ability = None
+    necrozma = battle.active_pokemon
+    groudon = battle.opponent_active_pokemon
+
+    necrozma.switch_out()
+    groudon.switch_in()
+    groudon._ability = None
 
     battle.parse_message(
         [
@@ -551,28 +556,94 @@ def test_battle_request_and_interactions(example_request):
             "[of] p1a: Groudon",
         ]
     )
-    assert battle.active_pokemon.ability == "waterabsorb"
-    battle.active_pokemon._ability = None
+    assert necrozma.ability == "waterabsorb"
 
-    battle.active_pokemon.item = GenData.UNKNOWN_ITEM
+    necrozma.item = GenData.UNKNOWN_ITEM
     battle.parse_message(
         ["", "-heal", "p2a: Necrozma", "200/265", "[from] item: Leftovers"]
     )
-    assert battle.active_pokemon.item == "leftovers"
-    battle.active_pokemon.item = None
+    assert necrozma.item == "leftovers"
+    necrozma.item = None
 
-    battle.opponent_active_pokemon.item = GenData.UNKNOWN_ITEM
+    groudon.item = GenData.UNKNOWN_ITEM
     battle.parse_message(
         ["", "-heal", "p1a: Groudon", "200/265", "[from] item: Leftovers"]
     )
-    assert battle.opponent_active_pokemon.item == "leftovers"
-    battle.opponent_active_pokemon.item = None
+    assert groudon.item == "leftovers"
+    groudon.item = None
 
-    battle.opponent_active_pokemon.item = None
+    groudon.item = None
     battle.parse_message(
         ["", "-heal", "p1a: Groudon", "200/265", "[from] item: Sitrus Berry"]
     )
-    assert battle.opponent_active_pokemon.item is None
+    assert groudon.item is None
+
+    # Test temporary types and abilities
+
+    necrozma._ability = "prismarmor"
+    groudon._ability = "desolateland"
+    battle.parse_message(["", "move", "p2a: Necrozma", "Worry Seed", "p1a: Groudon"])
+    assert groudon.ability == "desolateland"
+    battle.parse_message(
+        ["", "-endability", "p1a: Groudon", "Desolate Land", "[from] move: Worry Seed"]
+    )
+    assert groudon.ability is None
+    battle.parse_message(
+        ["", "-ability", "p1a: Groudon", "Insomnia", "[from] move: Worry Seed"]
+    )
+    assert groudon.ability == "insomnia"
+    groudon.switch_out()
+    groudon.switch_in()
+    assert groudon.ability == "desolateland"
+
+    battle.parse_message(["", "move", "p1a: Groudon", "Skill Swap", "p2a: Necrozma"])
+    battle.parse_message(
+        [
+            "",
+            "-activate",
+            "p1a: Groudon",
+            "move: Skill Swap",
+            "Prism Armor",
+            "Desolate Land",
+            "[of] p2a: Necrozma",
+        ]
+    )
+    assert groudon.ability == "prismarmor"
+    assert necrozma.ability == "desolateland"
+    groudon.switch_in()
+    groudon.switch_out()
+    assert groudon.ability == "desolateland"
+
+    battle.parse_message(["", "switch", "p1a: Ho-oh", "Ho-oh, L82", "100/100"])
+    hooh = battle.opponent_active_pokemon
+    assert hooh.type_1 == PokemonType.FIRE
+    assert hooh.type_2 == PokemonType.FLYING
+    assert hooh.types == [PokemonType.FIRE, PokemonType.FLYING]
+
+    battle.parse_message(["", "move", "p1a: Ho-oh", "Burn Up", "p2a: Necrozma"])
+    battle.parse_message(
+        ["", "-start", "p1a: Ho-oh", "typechange", "???/Flying", "[from] move: Burn Up"]
+    )
+
+    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_in()
+    assert hooh.type_1 == PokemonType.FIRE
+    assert hooh.type_2 == PokemonType.FLYING
+    assert hooh.types == [PokemonType.FIRE, PokemonType.FLYING]
+
+    battle.parse_message(["", "move", "p2a: Necrozma", "Soak", "p1a: Ho-oh"])
+    battle.parse_message(["", "-start", "p1a: Ho-oh", "typechange", "Water"])
+    assert hooh.type_1 == PokemonType.WATER
+    assert hooh.type_2 is None
+    assert hooh.types == [PokemonType.WATER]
+    hooh.switch_out()
+    hooh.switch_in()
+    assert hooh.type_1 == PokemonType.FIRE
+    assert hooh.type_2 == PokemonType.FLYING
+    assert hooh.types == [PokemonType.FIRE, PokemonType.FLYING]
 
 
 def test_end_illusion():
diff --git a/unit_tests/environment/test_pokemon.py b/unit_tests/environment/test_pokemon.py
index 4bc9aacdb..635c40cd7 100644
--- a/unit_tests/environment/test_pokemon.py
+++ b/unit_tests/environment/test_pokemon.py
@@ -307,3 +307,42 @@ def test_stats(example_request, showdown_format_teams):
         "spd": 120,
         "spe": 82,
     }
+
+
+def test_temporary():
+    furret = Pokemon(species="furret", gen=8)
+    furret._ability = "adaptability"
+
+    assert furret.types == [PokemonType.NORMAL]
+    assert furret.stab_multiplier == 2
+
+    furret.start_effect("typechange", "???/Fighting")
+    furret.start_effect("move: Skill Swap", ["Levitate", "Adaptability"])
+
+    assert furret.types == [PokemonType.THREE_QUESTION_MARKS, PokemonType.FIGHTING]
+    assert furret.ability == "levitate"
+    assert furret.damage_multiplier(PokemonType.PSYCHIC) == 1
+    assert furret.stab_multiplier == 1.5
+
+    furret.start_effect("typechange", "Fighting")
+    assert furret.damage_multiplier(PokemonType.PSYCHIC) == 2
+
+    furret.switch_out()
+    furret.switch_in()
+
+    assert furret.ability == "adaptability"
+
+    furret.terastallize("dragon")
+    assert furret.type_1 == PokemonType.DRAGON
+    assert furret.type_2 is None
+    assert furret.damage_multiplier(PokemonType.ICE) == 2
+
+    furret.set_temporary_ability("frisk")
+    assert furret.ability == "frisk"
+
+    furret.set_temporary_ability(None)
+    assert furret.ability is None
+
+    furret.switch_out()
+    assert furret.ability == "adaptability"
+    assert furret.types == [PokemonType.DRAGON]