Add a new init subcommand

To download resources.

lokrez/__init__.py

@@ -7,6 +7,7 @@ import os, os.path
 import sys
 
 import export
+import resources
 import smashgg
 import version
 
@@ -30,6 +31,29 @@ def main():
                         default = None,
                          help = "the authentication token to use" )
 
+    parser.add_argument(
+            "--rootdir", "-RD",
+            default = None,
+            help = "The directories containing this script, defaults to '.'",
+            )
+
+    # -------------------------------------------------------------------------
+    init_parser = subparsers.add_parser(
+            "init",
+            )
+
+    init_parser.add_argument(
+            "game",
+            default = "ssbu",
+            help = "The game you want to initialize the resources for",
+            )
+
+    init_parser.add_argument(
+            "--imgdir", "-ID",
+            default = "res/ssbu",
+            help = "The directory we should download the resources to",
+            )
+
     # -------------------------------------------------------------------------
     top8_parser = subparsers.add_parser(
             "top8",
@@ -40,11 +64,6 @@ def main():
             default = None,
             help = "The tournament slug or id",
             )
-    top8_parser.add_argument(
-            "--rootdir", "-RD",
-            default = None,
-            help = "The directories containing this script, defaults to '.'",
-            )
     top8_parser.add_argument(
             "--imgdir", "-ID",
             default = "res/ssbu",
@@ -122,11 +141,18 @@ def main():
         args.rootdir = "." # TODO compute script root?
 
     # -------------------------------------------------------------------------
-    if args.command not in [ "top8" ]:
+    if args.command not in [ "init", "top8" ]:
         parser.print_help()
         return 1
 
     # -------------------------------------------------------------------------
+    if args.command == "init":
+        resources.download_res_ssbu(
+                dstdir = args.imgdir,
+                log = log,
+                )
+        return 0
+    # -------------------------------------------------------------------------
     if args.command == "top8":
 
         tournament = None

lokrez/resources.py

@@ -0,0 +1,148 @@
+import io
+import os
+import sys
+import zipfile
+
+import requests
+
+import characters_ssbu
+
+# -----------------------------------------------------------------------------
+def download_file(url, with_progressbar = True):
+
+    r = requests.get(
+            url,
+            stream = with_progressbar,
+            )
+
+    if not with_progressbar:
+        return io.BytesIO(r.content)
+
+    total = r.headers.get("content-length")
+
+    f = io.BytesIO()
+
+    if total is not None:
+        downloaded = 0
+        total = int(total)
+        for data in r.iter_content(
+                chunk_size = max(int(total/1000), 1024*1024),
+                ):
+            f.write(data)
+            downloaded += len(data)
+            done = int(50*downloaded/total)
+            sys.stdout.write( "\r[{}{}] ({:02d}%)".format(
+                "█" * done,
+                " " * (50-done),
+                done*2,
+                ) )
+            sys.stdout.flush()
+        sys.stdout.write("\n")
+    else:
+        f = write(r.content)
+
+    return f
+
+# -----------------------------------------------------------------------------
+def download_res_ssbu(dstdir, log=None):
+    """Downloads SSBU resources from spriters and rename them according to
+    lokrez expectations"""
+
+    stock_icons_url = "https://www.spriters-resource.com/download/111395/"
+
+    # -------------------------------------------------------------------------
+    # Download stock icons
+    log.info("Downloading stock icons...")
+    fstocks = download_file(stock_icons_url)
+    zfstocks = zipfile.ZipFile(fstocks)
+
+    # Iter over each character
+    for character in characters_ssbu.EVERYONE:
+        log.info("Downloading images for {}...".format(character.name))
+
+        # Create directory for this character
+        chardir = os.path.join(
+                dstdir,
+                character.name,
+                )
+        try:
+            os.mkdir( chardir )
+
+        except FileExistsError:
+            log.warning(
+                    "Directory already exists for {}".format(character.name)
+                    )
+                
+            if os.listdir( chardir ):
+                log.warning(
+                        "Directory not empty for {}, skipping" \
+                                .format(character.name)
+                        )
+                continue
+
+        # Download urls & write image files
+        for url in character.spritersurls:
+            f = download_file(url)
+            with zipfile.ZipFile(f) as zfchar:
+                for zf in [zfchar,zfstocks]:
+                    for source_filename in zf.namelist():
+
+                        if "No Gamma Fix" in source_filename:
+                            continue
+
+                        if os.path.basename(source_filename) in ["","Tag.txt"]:
+                            continue
+
+                        if character.codename not in source_filename:
+                            continue
+
+                        target_filename = os.path.basename(source_filename)
+
+                        target_filename = target_filename.replace(
+                                character.codename,
+                                character.name,
+                                )
+
+                        log.debug("Writing file '{}'".format(target_filename))
+
+                        target_filename = os.path.join(
+                                chardir,
+                                target_filename,
+                                )
+
+                        with open(target_filename, "wb") as tf:
+                            tf.write(zf.read(source_filename))
+
+# =============================================================================
+if __name__ == '__main__':
+
+    import argparse
+    import logging
+    import tempfile
+
+    logging.basicConfig(
+            level = logging.DEBUG,
+            format = "%(message)s",
+            )
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+            "dstdir",
+            default = None,
+            help = "directory where to store the downloaded resources " \
+                   "(default to a temporary file)",
+                   )
+
+    args = parser.parse_args()
+
+    if args.dstdir is None:
+        args.dstdir = tempfile.mkdtemp()
+        logging.warning(
+                "Storing in temporary directory : {}".format(args.dstdir)
+                )
+
+    download_res_ssbu(
+            dstdir = args.dstdir,
+            log = logging,
+            )