Merge pull request #19 from hsahovic/fix-pokemon-data-bugs

Fix pokemon data bugs

.pre-commit-config.yaml

@@ -1,5 +1,5 @@
 default_language_version:
-  python: python3.7.4
+  python: python3.6.9
 repos:
   - repo: https://github.com/ambv/black
     rev: stable

setup.py

@@ -20,7 +20,7 @@
 EMAIL = "[email protected]"
 AUTHOR = "Haris Sahovic"
 REQUIRES_PYTHON = ">=3.6.0"
-VERSION = "0.0.5"
+VERSION = "0.0.6"
 
 # What packages are required for this module to be executed?
 with open("requirements.txt") as requirements:

src/poke_env/environment/battle.py

@@ -96,15 +96,20 @@ def __init__(self, battle_tag: str, username: str, logger: Logger):  # pyre-igno
         self._team: Dict[str, Pokemon] = {}
         self._opponent_team: Dict[str, Pokemon] = {}
 
-    def get_pokemon(self, identifier: str, force_self_team: bool = False) -> Pokemon:
+    def get_pokemon(
+        self, identifier: str, force_self_team: bool = False, details: str = ""
+    ) -> Pokemon:
         """Returns the Pokemon object corresponding to given identifier. Can force to
         return object from the player's team if force_self_team is True. If the Pokemon
-        object does not exist, it will be created.
+        object does not exist, it will be created. Details can be given, which is
+        necessary to initialize alternate forms (eg. alolan pokemons) properly.
 
         :param identifier: The identifier to use to retrieve the pokemon.
         :type identifier: str
         :param force_self_team: Wheter to force returning a Pokemon from the player's
             team. Defaults to False.
+        :type details: str, optional
+        :param details: Detailled information about the pokemon. Defaults to ''.
         :type force_self_team: bool, optional, defaults to False
         :return: The corresponding pokemon object.
         :rtype: Pokemon
@@ -114,7 +119,10 @@ def get_pokemon(self, identifier: str, force_self_team: bool = False) -> Pokemon
         if identifier[3] != " ":
             identifier = identifier[:2] + identifier[3:]
             species = identifier[5:]
-        species = identifier[4:]
+        if details:
+            species = details.split(", ")[0]
+        else:
+            species = identifier[4:]
 
         if is_mine or force_self_team:
             team: Dict[str, Pokemon] = self.team
@@ -133,13 +141,13 @@ def get_pokemon(self, identifier: str, force_self_team: bool = False) -> Pokemon
 
             return team[identifier]
 
-    def _end_illusion(self, pokemon_name: str):
+    def _end_illusion(self, pokemon_name: str, details: str):
         if pokemon_name[:2] == self._player_role:
             active = self.active_pokemon
         else:
             active = self.opponent_active_pokemon
 
-        pokemon = self.get_pokemon(pokemon_name)
+        pokemon = self.get_pokemon(pokemon_name, details=details)
         pokemon._set_hp(f"{active.current_hp}/{active.max_hp}")
         active._was_illusionned()
         pokemon._switch_in()
@@ -271,7 +279,7 @@ async def _parse_message(self, split_message: List[str]) -> None:
             self._in_team_preview = True
         elif split_message[1] in ["drag", "switch"]:
             pokemon, details, hp_status = split_message[2:5]
-            self._switch(pokemon, hp_status)
+            self._switch(pokemon, details, hp_status)
         elif split_message[1] == "faint":
             pokemon = split_message[2]
             self.get_pokemon(pokemon)._faint()
@@ -297,7 +305,8 @@ async def _parse_message(self, split_message: List[str]) -> None:
             self.register_pokemon(player, details, item)
         elif split_message[1] == "replace":
             pokemon = split_message[2]
-            self._end_illusion(pokemon)
+            details = split_message[3]
+            self._end_illusion(pokemon, details)
         elif split_message[1] == "rule":
             self._rules.append(split_message[2])
         elif split_message[1] == "start":
@@ -407,14 +416,14 @@ def _side_start(self, side, condition):
     def _swap(self, *args, **kwargs):
         self.logger.warning("swap method in Battle is not implemented")
 
-    def _switch(self, pokemon, hp_status):
+    def _switch(self, pokemon, details, hp_status):
         identifier = pokemon.split(":")[0][:2]
         if identifier == self._player_role:
             self.active_pokemon._switch_out()
         else:
             if self.opponent_team:
                 self.opponent_active_pokemon._switch_out()
-        pokemon = self.get_pokemon(pokemon)
+        pokemon = self.get_pokemon(pokemon, details=details)
         pokemon._switch_in()
         pokemon._set_hp_status(hp_status)
 
@@ -424,7 +433,7 @@ async def _tied(self):
     def _update_team_from_request(self, side: Dict) -> None:
         for pokemon in side["pokemon"]:
             self.get_pokemon(
-                pokemon["ident"], force_self_team=True
+                pokemon["ident"], force_self_team=True, details=pokemon["details"]
             )._update_from_request(pokemon)
 
     async def _won_by(self, player_name: str):

src/poke_env/environment/pokemon.py

@@ -208,7 +208,7 @@ def _update_from_pokedex(self, species: str) -> None:
         if len(dex_entry["types"]) == 1:
             self._type_2 = None
         else:
-            self._type_1 = PokemonType.from_name(dex_entry["types"][1])
+            self._type_2 = PokemonType.from_name(dex_entry["types"][1])
 
         self._possible_abilities = dex_entry["abilities"]
         self._heightm = dex_entry["heightm"]

unit_tests/environment/test_pokemon.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 from poke_env.environment.pokemon import Pokemon
+from poke_env.environment.pokemon_type import PokemonType
 
 
 def test_pokemon_moves():
@@ -13,3 +14,38 @@ def test_pokemon_moves():
 
     assert not mon.preparing
     assert not mon.must_recharge
+
+
+def test_pokemon_types():
+    # single type
+    mon = Pokemon(species="pikachu")
+    assert mon.type_1 == PokemonType.ELECTRIC
+    assert mon.type_2 is None
+
+    # dual type
+    mon = Pokemon(species="garchomp")
+    assert mon.type_1 == PokemonType.DRAGON
+    assert mon.type_2 == PokemonType.GROUND
+
+    # alolan forms
+    mon = Pokemon(species="raichualola")
+    assert mon.type_1 == PokemonType.ELECTRIC
+    assert mon.type_2 == PokemonType.PSYCHIC
+
+    # megas
+    mon = Pokemon(species="altaria")
+    assert mon.type_1 == PokemonType.DRAGON
+    assert mon.type_2 == PokemonType.FLYING
+
+    mon._mega_evolve("altariaite")
+    assert mon.type_1 == PokemonType.DRAGON
+    assert mon.type_2 == PokemonType.FAIRY
+
+    # primals
+    mon = Pokemon(species="groudon")
+    assert mon.type_1 == PokemonType.GROUND
+    assert mon.type_2 is None
+
+    mon._primal()
+    assert mon.type_1 == PokemonType.GROUND
+    assert mon.type_2 == PokemonType.FIRE