- WIP re-working evolution chain parsing

ui/workers/gather_evolutions_worker.py

@@ -8,6 +8,7 @@ from db import db
 import re
 
 from utility.functions import get_form_name, get_display_name, parse_pfic
+from utility.data import non_evolution_forms
 
 class GatherEvolutionsWorkerSignals(QObject):
     finished = pyqtSignal(list)
@@ -25,23 +26,27 @@ class GatherEvolutions(QRunnable):
         except Exception as e:
             print(f"Error gathering Pokémon home storage status: {e}")
 
-    def gather_evolution_data(self, force_refresh = False):
+    def gather_evolution_data(self, force_refresh = True):
         all_pokemon_forms = db.get_list_of_pokemon_forms()
-        evolutions = []
+        #evolutions = []
+        evolutions = {}
 
         for pokemon_form in all_pokemon_forms:
             print(f"Processing {get_display_name(pokemon_form)}'s evolutions")
             pokemon_name = pokemon_form["name"]
             form = get_form_name(pokemon_form)
 
+            if pokemon_form["form_name"] and any(s in pokemon_form["form_name"] for s in non_evolution_forms):
+                continue
+
             cache_record_name = f"chain_{pokemon_name}_{form}"
             if force_refresh:
                 cache.purge(cache_record_name)
 
-            cached_entry = cache.get(cache_record_name)
+            #cached_entry = cache.get(cache_record_name)
-            if cached_entry != None:
+            #if cached_entry != None:
-                evolutions.extend(cached_entry)
+            #    evolutions[pokemon_form["pfic"]] = cached_entry
-                continue
+            #    continue
 
             #form = get_form_name(pokemon_form, not pokemon_form["gender_relevant"])
             search_form = form
@@ -82,13 +87,17 @@ class GatherEvolutions(QRunnable):
                 continue
 
             evolution_chain = []
+            evolution_tree = None
             if pokemon_name == "Eevee":
-                evolution_chain = self.parse_eevee_evolution_chain(evolution_table, pokemon_form)
+                evolution_tree = self.parse_eevee_evolution_chain(evolution_table, pokemon_form)
                 #evolutions.append(evolution_chain)
             else:
-                evolution_chain = self.parse_evolution_chain(evolution_table, pokemon_form)
+                evolution_tree = self.parse_evolution_chain(evolution_table, pokemon_form)
                 #evolutions.append(evolution_chain)
 
+            if evolution_tree:
+                self.traverse_and_store(evolution_tree, evolutions, gender)
+
             chain = []
             for pokemon in evolution_chain:
                 from_pfic = self.get_pokemon_form_by_name(pokemon["pokemon"], pokemon["form"], gender=gender)
@@ -105,7 +114,7 @@ class GatherEvolutions(QRunnable):
                             "to_pfic": to_pfic,
                             "method": stage["method"] 
                         }
-                        evolutions.append(evolution_info)
+                        evolutions[pokemon_form["pfic"]] = evolution_info
                         chain.append(evolution_info)
                         
                         #insert_evolution_info(evolution_info)
@@ -121,7 +130,7 @@ class GatherEvolutions(QRunnable):
                             "to_pfic": to_pfic,
                             "method": branch["method"] 
                         }
-                        evolutions.append(evolution_info)
+                        evolutions[pokemon_form["pfic"]] = evolution_info
                         chain.append(evolution_info)
                         #EvolutionInfo(from_pfic, to_pfic, branch.method)
                         #insert_evolution_info(evolution_info)
@@ -132,90 +141,102 @@ class GatherEvolutions(QRunnable):
 
         return evolutions
     
+    def traverse_and_store(self, node, evolutions, gender):
+        """Helper function to traverse evolution tree and store evolutions."""
+        from_pfic = self.get_pokemon_form_by_name(node["pokemon"], node["form"], gender=gender)
+        if not from_pfic:
+            return
+
+        for next_stage in node["evolves_to"]:
+            to_pfic = self.get_pokemon_form_by_name(next_stage["pokemon"], next_stage["form"], gender=gender)
+            if to_pfic:
+                composite_key = f"{from_pfic}->{to_pfic}"
+                evolution_info = {
+                    "from_pfic": from_pfic,
+                    "to_pfic": to_pfic,
+                    "method": next_stage["method"]
+                }
+                evolutions[composite_key] = (evolution_info)
+                self.traverse_and_store(next_stage, evolutions, gender)
+    
     def parse_evolution_chain(self, table, pokemon_form, force_refresh = False):
-        cache_record_name = f"evo_{pokemon_form["pfic"]}"
+        cache_record_name = f"evo_{pokemon_form['pfic']}"
         if force_refresh:
             cache.purge(cache_record_name)
 
         cached_entry = cache.get(cache_record_name)
-        if cached_entry != None:
+        if cached_entry is not None:
             return cached_entry
 
-        main_chain = []
-        current_stage = None
-        pending_method = None
         form = get_form_name(pokemon_form, not pokemon_form["gender_relevant"])
 
         tbody = table.find('tbody', recursive=False)
         if not tbody:
-            return []
+            return None
 
         rows = tbody.find_all('tr', recursive=False)
         main_row = rows[0]
         branch_rows = rows[1:]
 
+        def create_stage(td):
+            pokemon_name = self.extract_pokemon_name(td)
+            evolution_form = self.extract_evolution_form(td, pokemon_name)
+            return {
+                "pokemon": pokemon_name,
+                "form": evolution_form,
+                "method": None,
+                "evolves_to": []
+            }
+
         # Parse main evolution chain
+        pending_method = None
+        root = None
+        current_stage = None
+
         for td in main_row.find_all('td', recursive=False):
             if td.find('table'):
-                # This TD contains Pokemon information
+                new_stage = create_stage(td)
-                pokemon_name = self.extract_pokemon_name(td)
+                new_stage["method"] = pending_method
-                stage = self.extract_stage_form(td)
-                evolution_form = self.extract_evolution_form(td, pokemon_name)
-                new_stage = {
-                        "pokemon":pokemon_name,
-                        "method": pending_method,
-                        "stage": stage,
-                        "form": evolution_form,
-                        "next_stage": None,
-                        "previous_stage": None,
-                        "branches": [],
-                        "pfic": pokemon_form["pfic"]
-                }
                 pending_method = None
+                if root is None:
+                    root = new_stage  # Assign the root node
                 if current_stage:
-                    current_stage["next_stage"] = new_stage
+                    current_stage["evolves_to"].append(new_stage)
-                    new_stage["previous_stage"] = current_stage  # Set the back link
                 current_stage = new_stage
-                main_chain.append(current_stage)
             else:
-                # This TD contains evolution method for the next Pokemon
                 pending_method = self.extract_evolution_method(td)
 
         # Parse branching evolutions
         for row in branch_rows:
-            branch_stage = None
             branch_method = None
+            branch_stage = None
+
             for td in row.find_all('td', recursive=False):
                 if td.find('table'):
-                    pokemon_name = self.extract_pokemon_name(td)
+                    new_stage = create_stage(td)
-                    stage = self.extract_stage_form(td)
+                    new_stage["method"] = branch_method
-                    evolution_form = self.extract_evolution_form(td, pokemon_name)
-                    new_stage = {
-                        "pokemon":pokemon_name,
-                        "method": branch_method,
-                        "stage": stage,
-                        "form": evolution_form,
-                        "next_stage": None,
-                        "previous_stage": None,
-                        "branches": [],
-                        "pfic": pokemon_form["pfic"]
-                    }
                     branch_method = None
                     if branch_stage:
-                        branch_stage["next_stage"] = new_stage
+                        branch_stage["evolves_to"].append(new_stage)
-                        new_stage["previous_stage"] = branch_stage  # Set the back link
                     branch_stage = new_stage
-                    # Find which main chain Pokemon this branches from
+
-                    for main_stage in main_chain:
+                    # Find which main chain Pokémon this branches from
-                        if td.get('rowspan') and main_stage.pokemon == pokemon_name:
+                    for main_stage in self.find_stages(root):
-                            main_stage["branches"].append(branch_stage)
+                        if td.get('rowspan') and main_stage["pokemon"] == new_stage["pokemon"]:
-                            branch_stage["previous_stage"] = main_stage  # Set the back link to the main chain
+                            main_stage["evolves_to"].append(branch_stage)
                             break
                 else:
                     branch_method = self.extract_evolution_method(td)
 
-        cache.set(cache_record_name, main_chain)
+        cache.set(cache_record_name, root)
-        return main_chain
+        return root
+    
+    def find_stages(self, node):
+        """Helper function to find all stages in the evolution chain recursively."""
+        stages = [node]
+        for stage in node["evolves_to"]:
+            stages.extend(self.find_stages(stage))
+        return stages
     
     def extract_pokemon_name(self, td: Tag) -> Optional[str]:
         name_tag = self.find_name_tag(td)
@@ -255,23 +276,34 @@ class GatherEvolutions(QRunnable):
         if not tbody:
             return []
         
+        def create_stage(td):
+            pokemon_name = self.extract_pokemon_name(td)
+            stage = self.extract_stage_form(td)
+            return {
+                "pokemon": pokemon_name,
+                "form": None,
+                "method": None,
+                "evolves_to": []
+            }
+
         rows = tbody.find_all('tr', recursive=False)
         eevee_row = rows[1]
         method_row = rows[2]
         eeveelutions_row = rows[3]
 
         eevee_td = eevee_row.find('td', recursive=False)
-        pokemon_name, stage = self.parse_pokemon_subtable(eevee_td)
+        eevee_stage = create_stage(eevee_td)
-        eevee_stage = {
+        #pokemon_name, stage = self.parse_pokemon_subtable(eevee_td)
-            "pokemon":pokemon_name,
+        #eevee_stage = {
-            "method": None,
+        #    "pokemon":pokemon_name,
-            "stage": stage,
+        #    "method": None,
-            "form": None,
+        #    "stage": stage,
-            "next_stage": None,
+        #    "form": None,
-            "previous_stage": None,
+        #    "next_stage": None,
-            "branches": [],
+        #    "previous_stage": None,
-            "pfic": pokemon_form["pfic"]
+        #    "branches": [],
-        }
+        #    "pfic": pokemon_form["pfic"]
+        #}
 
         methods = []
         for method in method_row.find_all('td', recursive=False):
@@ -280,24 +312,25 @@ class GatherEvolutions(QRunnable):
         eeveelutions = []
         index = 0
         for eeveelution in eeveelutions_row.find_all('td', recursive=False):
-            pokemon_name, stage = self.parse_pokemon_subtable(eeveelution)
+            #pokemon_name, stage = self.parse_pokemon_subtable(eeveelution)
-            eeveelution_stage = {
+            #eeveelution_stage = {
-                "pokemon":pokemon_name,
+            #    "pokemon":pokemon_name,
-                "method": methods[index],
+            #    "method": methods[index],
-                "stage": stage,
+            #    "stage": stage,
-                "form": None,
+            #    "form": None,
-                "next_stage": None,
+            #    "next_stage": None,
-                "previous_stage": None,
+            #    "previous_stage": None,
-                "branches": [],
+            #    "branches": [],
-                "pfic": pokemon_form["pfic"]
+            #    "pfic": pokemon_form["pfic"]
-            }
+            #}
-            eeveelution_stage["previous_stage"] = eevee_stage  # Set the back link to Eevee
+            eeveelution_stage = create_stage(eeveelution)
+            #eeveelution_stage["previous_stage"] = eevee_stage  # Set the back link to Eevee
             eeveelutions.append(eeveelution_stage)
             index += 1
 
-        eevee_stage["branches"] = eeveelutions  # Set the branches directly, not as a nested list
+        eevee_stage["evolves_to"] = eeveelutions  # Set the branches directly, not as a nested list
 
-        return [eevee_stage]
+        return eevee_stage
     
     def parse_pokemon_subtable(self, td):
         if td.find('table'):

utility/data.py

@@ -303,3 +303,9 @@ default_forms = [
     "Masterpiece Form",
     "Chest Form"
 ]
+
+non_evolution_forms = [
+    "Mega", 
+    "Dynamax",
+    "Gigantamax"
+]