diff --git a/lokrez/resources.py b/lokrez/resources.py index 12ca506..8e6fb3f 100644 --- a/lokrez/resources.py +++ b/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,86 +45,152 @@ def download_file(url, with_progressbar = True, proxy = None): return f # ----------------------------------------------------------------------------- -def download_res_ssbu(dstdir, proxy = None, log=None): - """Downloads SSBU resources from spriters and rename them according to - lokrez expectations""" +def download_res( + dstdir, + 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 - # ------------------------------------------------------------------------- - # Download stock icons - log.warning("Downloading stock icons...") - fstocks = download_file(stock_icons_url, proxy = proxy) - zfstocks = zipfile.ZipFile(fstocks) + # Select default source if needed + if not source: + if game.GAME.name == "pplus": + source = "smashlyon" + elif game.GAME.name == "ssbu": + source = "spriters" + + if source not in ["smashlyon", "spriters"]: + raise NotImplementedError( + "The only working sources are 'smashlyon' and 'spriters'", + ) + + gamedir = dstdir / game.GAME.name + + try: + 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 EVERYONE: + for character in game.EVERYONE: log.warning("Downloading images for {}...".format(character.name)) - # Create directory for this character - chardir = dstdir / character.name + chardir = gamedir / character.name - try: - chardir.mkdir() - - except FileExistsError: - log.info( - "Directory already exists for {}".format(character.name) - ) - + if not store_raw: try: - next(chardir.iterdir()) - log.warning( - "Directory not empty for {}, skipping" \ + chardir.mkdir(parents=True) + except FileExistsError: + log.info( + "Directory already exists for {}" \ .format(character.name) ) - continue - except StopIteration: - log.info( - "Directory empty, downloading", - ) + + try: + next(chardir.iterdir()) + log.warning( + "Directory not empty for {}, skipping" \ + .format(character.name) + ) + continue + except StopIteration: + log.info( + "Directory empty, downloading", + ) # Download urls & write image files - for url in character.spritersurls: - try: - f = download_file(url, proxy = proxy) - except Exception as e: + 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: - log.warning("Download failed ({}), retrying".format(e)) f = download_file(url, proxy = proxy) + except Exception as e: - log.error("Download failed({})".format(e)) - log.debug(e, exc_info = True) - continue + try: + # Try the download a 2nd time + log.warning("Download failed ({}), retrying".format(e)) + f = download_file(url, proxy = proxy) + except Exception as e: + log.error("Download failed({})".format(e)) + log.debug(e, exc_info = True) + continue - with zipfile.ZipFile(f) as zfchar: - for zf in [zfchar,zfstocks]: - for source_filename in zf.namelist(): + # We save the file in cache if it's the second time we need + # 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 "No Gamma Fix" in source_filename: - continue - if character.codename not in source_filename: - continue + # 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()) - target_filename = pathlib.Path(source_filename).name + # Add symlink for readablity + os.symlink( + str(outfile_name), + str( gamedir/ + "{charname}.{nb}.zip".format( + charname = character.name, + nb = url_nb+1, + )), + ) - if target_filename in ["","Tag.txt"]: - continue + continue - target_filename = pathlib.Path(source_filename).name + # otherwise: get the characters pictures and write them in the + # outdir + with zipfile.ZipFile(f) as zf: + for source_filename in zf.namelist(): - target_filename = target_filename.replace( - character.codename, - character.name, - ) + if "No Gamma Fix" in source_filename: + continue - log.debug("Writing file '{}'".format(target_filename)) + if character.codename not in source_filename: + continue - target_filename = chardir / target_filename + target_filename = pathlib.Path(source_filename).name - with open(str(target_filename), "wb") as tf: - tf.write(zf.read(source_filename)) + if target_filename in ["","Tag.txt"]: + continue + + target_filename = pathlib.Path(source_filename).name + + target_filename = target_filename.replace( + character.codename, + character.name, + ) + + log.debug("Writing file '{}'".format(target_filename)) + + target_filename = chardir / target_filename + + with open(str(target_filename), "wb") as tf: + tf.write(zf.read(source_filename)) # ============================================================================= if __name__ == '__main__':