Adding functionality to build cache (experimental)

Author: Shinga13

main.py

@@ -601,7 +601,7 @@ def launch():
             help='Provide this flag to build the cache instead of starting the app.')
         args, _ = argparser.parse_known_args()
         if args.build_cache:
-            exit_code = build_cache(Path(Launcher.base_path(), '.config'))
+            exit_code = build_cache(Path(Launcher.base_path()))
             sys.exit(exit_code)
         exit_code = SETS(
                 theme=Launcher.theme, args=args,

src/app.py

@@ -8,6 +8,7 @@
     ABOTTOM, ACENTER, AHCENTER, ALEFT, ARIGHT, ATOP, AVCENTER, CAREERS, FACTIONS, MARKS,
     PRIMARY_SPECS, RARITIES, SCROLLOFF, SCROLLON, SECONDARY_SPECS, SMAXMAX, SMAXMIN, SMINMAX,
     SMINMIN)
+from .datafunctions import cache_skills
 from .iofunc import (
         create_folder, delete_folder_contents, get_asset_path, load_icon, load_json, open_url,
         store_json)
@@ -31,7 +32,7 @@ class SETS():
             skill_unlock_callback, spec_combo_callback, species_combo_callback, switch_main_tab,
             tier_callback)
     from .datafunctions import (
-            autosave, backup_cargo_data, cache_skills, empty_build,
+            autosave, backup_cargo_data, empty_build,
             init_backend, load_legacy_build_image)
     from .export import get_build_markdown
     from .splash import enter_splash, exit_splash, splash_text
@@ -430,7 +431,7 @@ def setup_build_frames(self):
         """
         self.setup_space_build_frame()
         self.setup_ground_build_frame()
-        self.cache_skills()
+        cache_skills(self.cache.skills, self.app_dir)
         self.setup_space_skill_frame()
         self.setup_ground_skill_frame()
 

src/buildupdater.py

@@ -25,7 +25,7 @@ def load_build(self):
         self.widgets.ship['button'].setText(ship)
         ship_data = self.cache.ships[ship]
         exec_in_thread(
-                self, get_ship_image, self, ship_data['image'],
+                self, get_ship_image, self, ship_data['image'][5:],
                 result=lambda img: self.widgets.ship['image'].set_image(*img))
         tier = self.build['space']['tier']
         ship_tier = ship_data['tier']

src/callbacks.py

@@ -323,7 +323,7 @@ def select_ship(self):
     self.widgets.ship['button'].setText(new_ship)
     ship_data = self.cache.ships[new_ship]
     exec_in_thread(
-            self, get_ship_image, self, ship_data['image'],
+            self, get_ship_image, self, ship_data['image'][5:],
             result=lambda img: self.widgets.ship['image'].set_image(*img))
     tier = ship_data['tier']
     self.widgets.ship['tier'].clear()

src/datafunctions.py

@@ -11,6 +11,7 @@
 from requests.exceptions import (
         ConnectionError as requests__ConnectionError, Timeout as requests__Timeout)
 from requests_html import Element
+from urllib.parse import unquote_plus
 
 from .buildupdater import get_boff_spec, load_build, load_skill_pages
 from .constants import (
@@ -19,8 +20,9 @@
         STARSHIP_TRAIT_QUERY_URL, TRAIT_QUERY_URL, WIKI_IMAGE_URL)
 from .iofunc import (
         auto_backup_cargo_file, browse_path, cache_cargo_data, copy_file, download_image,
-        fetch_html, get_asset_path, get_cached_cargo_data, get_cargo_data, get_downloaded_images,
-        image, load_image, load_json, read_env_file, retrieve_image, store_json, store_to_cache)
+        download_images_fast, fetch_html, get_asset_path, get_cached_cargo_data, get_cargo_data,
+        get_downloaded_images, image, load_image, load_json, read_env_file, retrieve_image,
+        store_json, store_to_cache)
 from .splash import enter_splash, exit_splash, splash_text
 from .textedit import (
         create_equipment_tooltip, create_trait_tooltip, dewikify, parse_wikitext,
@@ -347,7 +349,8 @@ def download_images(self, threaded_worker: ThreadObject):
             no_retry_images.add(img)
         else:
             self.cache.images_failed.pop(img)
-    images = self.cache.images_set - no_retry_images - get_downloaded_images(self)
+    images = self.cache.images_set - no_retry_images - get_downloaded_images(
+        Path(self.config['config_subfolders']['images']))
     img_folder = self.config['config_subfolders']['images']
 
     images_to_download = images - self.cache.boff_abilities['all'].keys()
@@ -378,16 +381,16 @@ def cache_doff_single(self, cache: dict, doff: dict):
         cache[doff['spec']][doff['description']] = doff
 
 
-def cache_skills(self):
+def cache_skills(skill_cache: dict[str, dict], app_directory: str):
     """
     Loads skills into cache.
     """
-    space_skill_data = load_json(get_asset_path('space_skills.json', self.app_dir))
-    self.cache.skills['space'] = space_skill_data['space']
-    self.cache.skills['space_unlocks'] = space_skill_data['space_unlocks']
-    ground_skill_data = load_json(get_asset_path('ground_skills.json', self.app_dir))
-    self.cache.skills['ground'] = ground_skill_data['ground']
-    self.cache.skills['ground_unlocks'] = ground_skill_data['ground_unlocks']
+    space_skill_data = load_json(get_asset_path('space_skills.json', app_directory))
+    skill_cache['space'] = space_skill_data['space']
+    skill_cache['space_unlocks'] = space_skill_data['space_unlocks']
+    ground_skill_data = load_json(get_asset_path('ground_skills.json', app_directory))
+    skill_cache['ground'] = ground_skill_data['ground']
+    skill_cache['ground_unlocks'] = ground_skill_data['ground_unlocks']
 
 
 def autosave(self):
@@ -1053,14 +1056,74 @@ class ThisIsTerribleError(RuntimeError):
         store_json(self.cache.boff_abilities, filepath)
 
 
-def build_cache(config_path: Path) -> int:
+def get_icon_set(cargo_dir: Path) -> set[str]:
+    """
+    Creates set of all required icons from cargo data and required static images.
+
+    Parameters:
+    - :param cargo_dir: path to cargo data directory
+    """
+    images_set = set()
+    equipment_cargo_data = load_json(str(cargo_dir / 'equipment.json'))
+    equipment_types = set(EQUIPMENT_TYPES.keys())
+    elite_hangar = {
+        'Hangar - Elite Federation Mission Scout Ships',
+        'Hangar - Elite Valor Fighters'
+    }
+    for item in equipment_cargo_data:
+        if item['type'] in equipment_types:
+            if item['type'] == 'Hangar Bay' and item['name'] not in elite_hangar and (
+                    item['name'].startswith('Hangar - Advanced')
+                    or item['name'].startswith('Hangar - Elite')):
+                continue
+            images_set.add(sanitize_equipment_name(item['name']))
+    trait_cargo_data = load_json(str(cargo_dir / 'traits.json'))
+    for trait in trait_cargo_data:
+        if trait['type'] != 'doff' and trait['type'] != 'boff' and trait['name'] is not None:
+            if trait['icon_name'] is None:
+                images_set.add(trait['name'])
+            else:
+                images_set.add(trait['icon_name'])
+    shiptrait_cargo_data = load_json(str(cargo_dir / 'starship_traits.json'))
+    images_set |= set(ship_trait['name'] for ship_trait in shiptrait_cargo_data)
+    return images_set
+
+
+def get_skill_icons(skill_cache: dict[str, dict]) -> set[str]:
+    """
+    """
+    icons = set()
+    for rank_group in skill_cache['space']:
+        for skill_group in rank_group:
+            for skill_node in skill_group['nodes']:
+                icons.add(skill_node['image'])
+    for skill_group in skill_cache['ground']:
+        for skill_node in skill_group['nodes']:
+            icons.add(skill_node['image'])
+    return icons
+
+
+def get_boff_icons(boff_cache: dict[str, dict]) -> set[str]:
+    """
+    """
+    return set(boff_cache['all'].keys())
+
+
+def get_ship_icons(ship_list: list[dict[str]]) -> set[str]:
+    """
+    """
+    return set(ship['image'][5:] for ship in ship_list)
+
+
+def build_cache(app_dir: Path) -> int:
     """
     Builds cache in config folder indicated by `config_path`. Returns status: success: `0`,
     failure: `1`
 
     Parameters:
     - :param config_path: path to build cache into
     """
+    config_path = app_dir / '.config'
     env_variables = read_env_file(config_path / '.env', ['SETS_CF_CLEARANCE', 'SETS_USER_AGENT'])
     requests_session = Session()
     if 'SETS_CF_CLEARANCE' in env_variables:
@@ -1077,6 +1140,32 @@ def build_cache(config_path: Path) -> int:
         cargo_dir / 'starship_traits.json', STARSHIP_TRAIT_QUERY_URL, requests_session))
     success.append(cache_cargo_data(cargo_dir / 'modifiers.json', MODIFIER_QUERY, requests_session))
     success.append(cache_cargo_data(cargo_dir / 'doffs.json', DOFF_QUERY_URL, requests_session))
+
+    image_dir = config_path / 'images'
+    downloaded_images = get_downloaded_images(image_dir)
+    ultimate_icons = {'Focused Frenzy', 'Probability Manipulation', 'EPS Corruption'}
+    images_set = (get_icon_set(cargo_dir) | ultimate_icons) - downloaded_images
+    if len(images_set) > 0:
+        download_images_fast(list(images_set), env_variables, image_dir)
+    skill_cache = dict()
+    cache_skills(skill_cache, app_dir)
+    skill_images = get_skill_icons(skill_cache) - downloaded_images
+    if len(skill_images) > 0:
+        download_images_fast(list(skill_images), env_variables, image_dir, image_suffix='.png')
+    boff_cache = load_json(cargo_dir / 'boff_abilities.json')
+    boff_images = get_boff_icons(boff_cache) - downloaded_images
+    if len(boff_images) > 0:
+        download_images_fast(
+            list(boff_images), env_variables, image_dir, image_suffix='_icon_(Federation).png')
+
+    downloaded_ship_images = set(
+        map(lambda x: unquote_plus(x), os.listdir(str(config_path / 'ship_images'))))
+    ship_list = load_json(str(cargo_dir / 'ship_list.json'))
+    ship_images = get_ship_icons(ship_list) - downloaded_ship_images
+    if len(ship_images) > 0:
+        download_images_fast(
+            list(ship_images), env_variables, config_path / 'ship_images', image_suffix='')
+
     if False in success:
         return 1
     return 0

src/iofunc.py

@@ -4,18 +4,34 @@
 from pathlib import Path
 from shutil import copyfile as shutil__copyfile, rmtree as shutil__rmtree
 import sys
+from threading import Thread
 from urllib.parse import quote_plus, unquote_plus
 from webbrowser import open as webbrowser_open
 
 from PySide6.QtGui import QIcon, QImage
 from PySide6.QtWidgets import QFileDialog
 import requests
+from requests.cookies import create_cookie as requests__create_cookie
 from requests_html import HTMLSession
 
 from .constants import WIKI_IMAGE_URL, WIKI_URL
 from .textedit import compensate_json
 
 
+class ReturnValueThread(Thread):
+    def __init__(self, target, args: tuple = tuple()):
+        super().__init__(target=target, args=args)
+        self._return = None
+
+    def run(self):
+        if self._target is not None:
+            self._return = self._target(*self._args)
+
+    def join(self):
+        super().join()
+        return self._return
+
+
 def browse_path(self, default_path: str = None, types: str = 'Any File (*.*)', save=False) -> str:
     """
     Opens file dialog prompting the user to select a file.
@@ -233,14 +249,6 @@ def alt_image(self, image_name: str, image_suffix: str) -> QImage:
         return image(self, image_name)
 
 
-def get_downloaded_images(self) -> set:
-    """
-    Returns set containing all images currently in the images folder.
-    """
-    img_folder = self.config['config_subfolders']['images']
-    return set(map(lambda x: unquote_plus(x)[:-4], os.listdir(img_folder)))
-
-
 def auto_backup_cargo_file(self, filename: str):
     """
     Backs up given cargo data file to the auto backups folder
@@ -259,6 +267,13 @@ def auto_backup_cargo_file(self, filename: str):
 # --------------------------------------------------------------------------------------------------
 
 
+def get_downloaded_images(images_dir: Path) -> set:
+    """
+    Returns set containing all images currently in the images folder.
+    """
+    return set(map(lambda x: unquote_plus(x)[:-4], os.listdir(str(images_dir))))
+
+
 def create_folder(path_to_folder):
     """
     Creates the folder at path_to_folder in case it does not exist.
@@ -499,3 +514,70 @@ def cache_cargo_data(cache_file: Path, url: str, session: requests.Session) -> b
             sys.stdout.write(
                 f'[Error] Decoding the response failed for the following URL:\n[Error] {url}\n')
     return False
+
+
+def download_image_session(
+        session: requests.Session, name: str, image_folder_path: Path,
+        failed_images: dict[str, int], image_suffix: str = '_icon.png'):
+    """
+    """
+    if image_suffix == '':
+        # exception for ship images
+        filepath = image_folder_path / quote_plus(name)
+        image_type = None
+    else:
+        filepath = image_folder_path / get_image_file_name(name)
+        image_type = 'png'
+    image_url = WIKI_IMAGE_URL + name.replace(' ', '_') + image_suffix
+    image_response = session.get(image_url)
+    image = QImage()
+    if image_response.ok:
+        image.loadFromData(image_response.content, image_type)
+        image.save(str(filepath))
+    else:
+        failed_images[name] = int(datetime.now().timestamp())
+
+
+def download_images_list(
+        images_list: list[str], env_variables: dict[str, str], images_path: Path,
+        image_suffix: str = '_icon.png') -> dict[str, int]:
+    """
+    """
+    requests_session = requests.Session()
+    if 'SETS_CF_CLEARANCE' in env_variables:
+        requests_session.cookies.set_cookie(
+            requests__create_cookie(name='cf_clearance', value=env_variables['SETS_CF_CLEARANCE']))
+    if 'SETS_USER_AGENT' in env_variables:
+        requests_session.headers['User-Agent'] = env_variables['SETS_USER_AGENT']
+    failed_images = dict()
+    for image_name in images_list:
+        download_image_session(
+            requests_session, image_name, images_path, failed_images, image_suffix)
+    return failed_images
+
+
+def download_images_fast(
+        images_list: list[str], env_variables: dict[str, str], images_dir: Path,
+        image_suffix: str = '_icon.png'):
+    """
+    Downloads images using multiple threads.
+    """
+    total_threads = 16
+    image_chunk_size = len(images_list) // total_threads
+    while image_chunk_size < 4 and total_threads > 1:
+        total_threads -= 1
+        image_chunk_size = len(images_list) // total_threads
+    threads: list[ReturnValueThread] = list()
+    for thread_num in range(total_threads):
+        if thread_num == total_threads - 1:
+            images = images_list[image_chunk_size * thread_num:]
+        else:
+            images = images_list[image_chunk_size * thread_num:image_chunk_size * (thread_num + 1)]
+        thread = ReturnValueThread(
+            target=download_images_list, args=(images, env_variables, images_dir, image_suffix))
+        thread.start()
+        threads.append(thread)
+    failed_images = dict()
+    for thread in threads:
+        failed_images.update(thread.join())
+    print(failed_images)