Update the resdl function for more than just ssbu

lokrez/resources.py

@@ -1,12 +1,12 @@
 import io
+import os
 import pathlib
 import sys
+import urllib
 import zipfile
 
 import requests
 
-from .characters_ssbu import EVERYONE
-
 # -----------------------------------------------------------------------------
 def download_file(url, with_progressbar = True, proxy = None):
 
@@ -45,31 +45,56 @@ def download_file(url, with_progressbar = True, proxy = None):
     return f
 
 # -----------------------------------------------------------------------------
-def download_res_ssbu(dstdir, proxy = None, log=None):
+def download_res(
-    """Downloads SSBU resources from spriters and rename them according to
+        dstdir,
-    lokrez expectations"""
+        game = None,
+        source = None,
+        store_raw = False,
+        proxy = None,
+        log = None,
+        ):
+    """TODO: Docstring for download_res_pplus.
+    :returns: TODO
 
-    stock_icons_url = "https://www.spriters-resource.com/download/111395/"
+    """
+    if not game:
+        return
 
-    # -------------------------------------------------------------------------
+    # Select default source if needed
-    # Download stock icons
+    if not source:
-    log.warning("Downloading stock icons...")
+        if game.GAME.name == "pplus":
-    fstocks = download_file(stock_icons_url, proxy = proxy)
+            source = "smashlyon"
-    zfstocks = zipfile.ZipFile(fstocks)
+        elif game.GAME.name == "ssbu":
+            source = "spriters"
 
-    # Iter over each character
+    if source not in ["smashlyon", "spriters"]:
-    for character in EVERYONE:
+        raise NotImplementedError(
-        log.warning("Downloading images for {}...".format(character.name))
+                "The only working sources are 'smashlyon' and 'spriters'",
+                )
 
-        # Create directory for this character
+    gamedir = dstdir / game.GAME.name
-        chardir = dstdir / character.name
 
     try:
-            chardir.mkdir()
+        gamedir.mkdir(parents=True)
+    except FileExistsError:
+        log.debug("Game directory already exist")
 
+    # A cache to save time
+    cache = {}
+
+    # Iter over each character
+    for character in game.EVERYONE:
+        log.warning("Downloading images for {}...".format(character.name))
+
+        chardir = gamedir /  character.name
+
+        if not store_raw:
+            try:
+                chardir.mkdir(parents=True)
             except FileExistsError:
                 log.info(
-                    "Directory already exists for {}".format(character.name)
+                        "Directory already exists for {}" \
+                                .format(character.name)
                         )
 
                 try:
@@ -85,11 +110,20 @@ def download_res_ssbu(dstdir, proxy = None, log=None):
                             )
 
         # Download urls & write image files
-        for url in character.spritersurls:
+        for url_nb, url in enumerate(character.res_urls[source]):
+
+            # If we have the file in cache, just get it
+            if url in cache and cache[url] is not None:
+                log.debug("Found url '{}' in cache".format(url))
+                f = cache[url]
+
+            else:
                 try:
                     f = download_file(url, proxy = proxy)
+
                 except Exception as e:
                     try:
+                        # Try the download a 2nd time
                         log.warning("Download failed ({}), retrying".format(e))
                         f = download_file(url, proxy = proxy)
                     except Exception as e:
@@ -97,8 +131,40 @@ def download_res_ssbu(dstdir, proxy = None, log=None):
                         log.debug(e, exc_info = True)
                         continue
 
-            with zipfile.ZipFile(f) as zfchar:
+                # We save the file in cache if it's the second time we need
-                for zf in [zfchar,zfstocks]:
+                # to download it.
+                if url in cache:
+                    log.debug("Saving url '{}' in cache".format(url))
+                    cache[url] = f
+                else:
+                    log.debug("Marking url '{}' in cache".format(url))
+                    cache[url] = None
+
+
+            # if store_raw: we just save the raw zip file
+            if store_raw:
+                outfile_name = pathlib.Path(
+                        urllib.parse.urlparse(url).path
+                        ) \
+                                .name
+                with open(str(gamedir/outfile_name), "wb") as outfile:
+                    outfile.write(f.getbuffer())
+
+                # Add symlink for readablity
+                os.symlink(
+                        str(outfile_name),
+                        str( gamedir/
+                        "{charname}.{nb}.zip".format(
+                            charname = character.name,
+                            nb = url_nb+1,
+                            )),
+                        )
+
+                continue
+
+            # otherwise: get the characters pictures and write them in the
+            # outdir
+            with zipfile.ZipFile(f) as zf:
                 for source_filename in zf.namelist():
 
                     if "No Gamma Fix" in source_filename: