diff --git a/ui/workers/gather_evolutions_worker.py b/ui/workers/gather_evolutions_worker.py index fad179e..712e424 100644 --- a/ui/workers/gather_evolutions_worker.py +++ b/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) - if cached_entry != None: - evolutions.extend(cached_entry) - continue + #cached_entry = cache.get(cache_record_name) + #if cached_entry != None: + # evolutions[pokemon_form["pfic"]] = cached_entry + # 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 - pokemon_name = self.extract_pokemon_name(td) - 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"] - } + new_stage = create_stage(td) + new_stage["method"] = pending_method pending_method = None + if root is None: + root = new_stage # Assign the root node if current_stage: - current_stage["next_stage"] = new_stage - new_stage["previous_stage"] = current_stage # Set the back link + current_stage["evolves_to"].append(new_stage) 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) - stage = self.extract_stage_form(td) - 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"] - } + new_stage = create_stage(td) + new_stage["method"] = branch_method branch_method = None if branch_stage: - branch_stage["next_stage"] = new_stage - new_stage["previous_stage"] = branch_stage # Set the back link + branch_stage["evolves_to"].append(new_stage) branch_stage = new_stage - # Find which main chain Pokemon this branches from - for main_stage in main_chain: - if td.get('rowspan') and main_stage.pokemon == pokemon_name: - main_stage["branches"].append(branch_stage) - branch_stage["previous_stage"] = main_stage # Set the back link to the main chain + + # Find which main chain Pokémon this branches from + for main_stage in self.find_stages(root): + if td.get('rowspan') and main_stage["pokemon"] == new_stage["pokemon"]: + main_stage["evolves_to"].append(branch_stage) break else: branch_method = self.extract_evolution_method(td) - cache.set(cache_record_name, main_chain) - return main_chain + cache.set(cache_record_name, root) + 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) @@ -254,6 +275,16 @@ class GatherEvolutions(QRunnable): tbody = table.find('tbody', recursive=False) 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] @@ -261,17 +292,18 @@ class GatherEvolutions(QRunnable): eeveelutions_row = rows[3] eevee_td = eevee_row.find('td', recursive=False) - pokemon_name, stage = self.parse_pokemon_subtable(eevee_td) - eevee_stage = { - "pokemon":pokemon_name, - "method": None, - "stage": stage, - "form": None, - "next_stage": None, - "previous_stage": None, - "branches": [], - "pfic": pokemon_form["pfic"] - } + eevee_stage = create_stage(eevee_td) + #pokemon_name, stage = self.parse_pokemon_subtable(eevee_td) + #eevee_stage = { + # "pokemon":pokemon_name, + # "method": None, + # "stage": stage, + # "form": None, + # "next_stage": None, + # "previous_stage": None, + # "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) - eeveelution_stage = { - "pokemon":pokemon_name, - "method": methods[index], - "stage": stage, - "form": None, - "next_stage": None, - "previous_stage": None, - "branches": [], - "pfic": pokemon_form["pfic"] - } - eeveelution_stage["previous_stage"] = eevee_stage # Set the back link to Eevee + #pokemon_name, stage = self.parse_pokemon_subtable(eeveelution) + #eeveelution_stage = { + # "pokemon":pokemon_name, + # "method": methods[index], + # "stage": stage, + # "form": None, + # "next_stage": None, + # "previous_stage": None, + # "branches": [], + # "pfic": pokemon_form["pfic"] + #} + 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'): diff --git a/utility/data.py b/utility/data.py index b53f354..f15e14c 100644 --- a/utility/data.py +++ b/utility/data.py @@ -302,4 +302,10 @@ default_forms = [ "Phony Form", "Masterpiece Form", "Chest Form" +] + +non_evolution_forms = [ + "Mega", + "Dynamax", + "Gigantamax" ] \ No newline at end of file