KAGARI: Add option to download multiple files at once

utils/scripts/kagari.py

@@ -25,8 +25,9 @@ def create_parser():
     usage = "kagari [-h] (dl [dl options] | search [search options])"
 
     desc = """Client to interact with the Kurisu API. Two main modes:
-    -dl: downloads the specified file, or a tar file of the entire database
-    -search: searches the database with filters for info on the files"""
+    -search: searches the database with filters for info on the files
+    -dl: downloads the target file with an ID, or multiple files with filters
+         To download the entire database, use neither ID nor filters"""
 
     parser = argparse.ArgumentParser(
         usage=usage,
@@ -52,7 +53,7 @@ def create_parser():
         "-i",
         "--id",
         help='In "dl" mode, if specified, downloads target video file.'
-        + "\nIf absent in this mode, downloads the tar of the whole base."
+        + "\nIf absent in this mode, filter arguments are used instead."
         + '\nIn "search" mode, returns info of the target file.\n ',
     )
 
@@ -70,32 +71,32 @@ def create_parser():
         "--check",
         action="store_true",
         help='In "dl" mode, first displays info on the targeted'
-        + "\nfile, and asks for confirmation before downloading.\n ",
+        + "\nfile(s), and asks for confirmation before downloading.\n ",
     )
 
     parser.add_argument(
         "-a",
         "--author",
-        help='In "search" mode, filters by author. (case insensitive)'
+        help='In either mode, filters by author. (case insensitive)',
     )
 
     parser.add_argument(
         "-c",
         "--category",
         dest="cat",
-        help='In "search" mode, filters by category.\n ',
+        help='In either mode, filters by category.\n ',
     )
 
     parser.add_argument(
-        "-t", "--type", help='In "search" mode, filters by type.\n '
+        "-t", "--type", help='In either mode, filters by type.\n '
     )
 
     parser.add_argument(
         "-s",
         "--string",
         dest="search",
-        help='In "search" mode, filters on the names of sources and songs'
-        + "\nmatching the given string. (case insensitive)\n "
+        help='In either mode, filters on the names of sources and songs'
+        + "\nmatching the given string. (case insensitive)\n ",
     )
 
     parser.add_argument(
@@ -110,7 +111,7 @@ def create_parser():
         "-j",
         "--json",
         action="store_true",
-        help='In "search" mode, output the result as a json.'
+        help='In "search" mode, output the result as a json.',
     )
 
     return parser
@@ -132,6 +133,44 @@ def make_request(url):
     return response
 
 
+def search_query(args):
+    "Searches the database with ID or filters, and returns the response json"
+    url = "https://kurisu.iiens.net/api"
+    url_values = dict()
+
+    # If ID, it's the only parameter needed
+    if args.id:
+        url_values["id"] = args.id
+    # No ID, building the query parameters
+    else:
+        for arg_name in ["author", "cat", "type", "search", "random"]:
+            arg = getattr(args, arg_name)
+            if arg:
+                url_values[arg_name] = arg
+
+    url_args = urlencode(url_values)
+    if url_args:
+        url += "?" + url_args
+
+    response = make_request(url)
+
+    try:
+        kara_json = json.loads(response.read().decode("utf-8"))
+    except ValueError:
+        emsg = "The response JSON could not be loaded."
+        print(emsg)
+        sys.exit(1)
+
+    # Handling an error json
+    if "error" in kara_json:
+        default_emsg = "An unknown error happened when searching the database."
+        emsg = kara_json.get("message", default_emsg)
+        print(emsg)
+        sys.exit(1)
+
+    return kara_json
+
+
 def build_kara_name(data):
     "Builds the name of the file from the response dictionary"
     src_name = data["source_name"]
@@ -148,9 +187,13 @@ def build_kara_name(data):
     return kara_name, kara_additional
 
 
-def kara_prompt(kara_string):
+def kara_prompt(kara_string="", kara_nbr=0):
     "Simple yes/no prompt to confirm the download or not"
-    prompt_str = "You are about to download " + kara_string
+    prompt_str = "You are about to download "
+    if kara_string:
+        prompt_str += kara_string
+    else:
+        prompt_str += f"{kara_nbr} files"
     prompt_str += "\nDo you want to proceed ? ([y]/n) "
 
     answers = {"y": True, "n": False, "ye": True, "yes": True, "no": False}
@@ -162,49 +205,38 @@ def kara_prompt(kara_string):
                 ans = "y"
         except KeyboardInterrupt:
             ans = "n"
-            print("\n")
+            print("")
         except EOFError:
             ans = "y"
 
     return answers[ans]
 
 
-def dl_mode(args):
-    "Downloads either the whole database as a tar, or a single video file"
+def dl_mode(args, kara_json):
+    "Downloads the targeted files"
 
-    # No ID, download the tar of the database
-    if not args.id:
-        url = "https://kurisu.iiens.net/kara.tar"
-        path = args.path if args.path else os.getcwd()
-        filename = path + "/kara.tar"
+    kara_id = None
+    kara_string = ""
+    kara_nbr = 0
 
-        print(f'Downloading and writing archive as "{filename}"')
-        response = make_request(url)
-
-        # Writing the tar file
-        with open(filename, 'wb') as out_file:
-            try:
-                shutil.copyfileobj(response, out_file)
-            except KeyboardInterrupt:
-                print("\nDownload interrupted by user input.")
-            else:
-                print("All done !")
-    # An ID, download target file
+    if len(kara_json) == 1:
+        if type(kara_json) is list:
+            kara_json = kara_json[0]
+        kara_name, kara_additional = build_kara_name(kara_json)
+        kara_string = '"' + kara_name + '"' + kara_additional
+        kara_id = kara_json["id"]
     else:
-        search_url = f"https://kurisu.iiens.net/api?id={args.id}"
-        search = make_request(search_url)
-        kara_data = json.loads(search.read().decode("utf-8"))
-        kara_name, kara_additional = build_kara_name(kara_data)
+        kara_nbr = len(kara_json)
 
     # If asked for check, display the prompt
     if args.check:
-            kara_string = '"' + kara_name + '"' + kara_additional
-            prompt = kara_prompt(kara_string)
+        prompt = kara_prompt(kara_string, kara_nbr)
         if not prompt:
             print("Download cancelled.")
             return
 
-        url = f"https://kurisu.iiens.net/api/download/{args.id}"
+    if kara_id:
+        url = f"https://kurisu.iiens.net/api/download/{kara_id}"
         path = args.path if args.path else os.getcwd()
         filename = path + "/" + kara_name + ".mkv"
 
@@ -212,71 +244,64 @@ def dl_mode(args):
         response = make_request(url)
 
         # Writing the video file
-        with open(filename, 'wb') as out_file:
-            try:
+        with open(filename, "wb") as out_file:
             shutil.copyfileobj(response, out_file)
-            except KeyboardInterrupt:
-                print("\nDownload interrupted by user input.")
-            else:
-                print("All done !")
-
-
-def search_mode(args):
-    "Performs a search request on the database and returns the answer"
-    url = "https://kurisu.iiens.net/api"
-    url_values = dict()
-
-    # If ID, it's the only parameter needed
-    if args.id:
-        url_values["id"] = args.id
-    # No ID, building the query parameters
     else:
-        for arg_name in ["author", "cat", "type", "search", "random"]:
-            arg = getattr(args, arg_name)
-            if arg:
-                url_values[arg_name] = arg
-
-    url_args = urlencode(url_values)
-    if url_args:
-        url += "?" + url_args
+        print(f"Download in progress: 0/{kara_nbr}", end="")
+        sys.stdout.flush()
+        for curr_nbr, kara in enumerate(kara_json):
+            kara_id = kara["id"]
+            kara_name, _ = build_kara_name(kara)
+            url = f"https://kurisu.iiens.net/api/download/{kara_id}"
+            path = args.path if args.path else os.getcwd()
+            filename = path + "/" + kara_name + ".mkv"
 
             response = make_request(url)
 
-    try:
-        kara_json = json.loads(response.read().decode("utf-8"))
-    except ValueError:
-        emsg = "The response JSON could not be loaded."
-        print(emsg)
-        sys.exit(1)
+            # Writing the video file
+            with open(filename, "wb") as out_file:
+                shutil.copyfileobj(response, out_file)
+            print(f"\rDownload in progress: {curr_nbr + 1}/{kara_nbr}", end="")
+            sys.stdout.flush()
+        print("")
+    print("All done !")
+
 
+def search_mode(args, kara_json):
+    "Performs a search request on the database and returns the answer"
     if args.json:
         pretty_output = json.dumps(kara_json, indent=4, ensure_ascii=False)
+        pretty_output += "\n"
     else:
         pretty_output = ""
         for res in kara_json:
-            new_kara      = (" NEW!" if res['is_new'] == "1" else "")
-            output_string = build_kara_name(res)
-            output_string = f"[{res['id']}] {output_string[0]}\n" + \
-                            f"{output_string[1]} in {res['author_year']}{new_kara}\n" + \
-                            f" Size: {res['size']} Kio\n"
+            new_kara = " NEW!" if res["is_new"] == "1" else ""
+            output_info = build_kara_name(res)
+            output_string = (
+                f"[{res['id']}] {output_info[0]}\n"
+                + f"{output_info[1]} in {res['author_year']}{new_kara}\n"
+                + f" Size: {res['size']} Kio\n"
+            )
             pretty_output += output_string
-            if res['popularity']:
-                pretty_output += " Popularity: " + res['popularity'] + "\n"
-            if res['upload_comment']:
-                for comment_line in wrap(res['upload_comment'], fix_sentence_endings=True):
+            if res["popularity"]:
+                pretty_output += " Popularity: " + res["popularity"] + "\n"
+            if res["upload_comment"]:
+                for comment_line in wrap(
+                    res["upload_comment"], fix_sentence_endings=True
+                ):
                     pretty_output += " " + comment_line + "\n"
             pretty_output += "\n"
 
     # Displaying the received JSON, or writing it to a file
     if not args.path:
-        print(pretty_output, end='')
+        print(pretty_output, end="")
     else:
         try:
             with open(args.path, "w") as output_file:
                 output_file.write(pretty_output)
         except OSError as err:
             emsg = "Cannot create/open the search output file: "
-            emsg += err
+            emsg += str(err)
             print(emsg)
             sys.exit(1)
 
@@ -286,10 +311,15 @@ def run():
     parser = create_parser()
     args = parser.parse_args()
 
+    kara_json = search_query(args)
+
     if args.mode == "dl":
-        dl_mode(args)
+        try:
+            dl_mode(args, kara_json)
+        except KeyboardInterrupt:
+            print("\nDownload interrupted by user input.")
     elif args.mode == "search":
-        search_mode(args)
+        search_mode(args, kara_json)
     else:
         print(f'Unknown mode "{args.mode}". Use either "dl" or "search".')
         sys.exit(1)