From 9f7187bb67ab18980fff6bedc286087a8ba1aa8c Mon Sep 17 00:00:00 2001 From: nichind Date: Mon, 9 Dec 2024 00:47:47 +0300 Subject: [PATCH] ref:proper cli args; download_callback --- pybalt/__init__.py | 11 ++++++++++- pybalt/__main__.py | 49 ++++++++++++++++++---------------------------- pybalt/cobalt.py | 11 +++++++++++ setup.py | 1 - 4 files changed, 40 insertions(+), 32 deletions(-) diff --git a/pybalt/__init__.py b/pybalt/__init__.py index a2a5528..bb001b5 100644 --- a/pybalt/__init__.py +++ b/pybalt/__init__.py @@ -1 +1,10 @@ -from .cobalt import Cobalt, Pybalt, download, get, check_updates, File, DownloadedFile, tl +from .cobalt import ( + Cobalt, + Pybalt, + download, + get, + check_updates, + File, + DownloadedFile, + tl, +) diff --git a/pybalt/__main__.py b/pybalt/__main__.py index b4352e2..0cdbfeb 100644 --- a/pybalt/__main__.py +++ b/pybalt/__main__.py @@ -9,79 +9,81 @@ async def _(): parser = argparse.ArgumentParser() parser.add_argument("url_arg", nargs="?", type=str, help=tl("URL_ARGUMENT")) - parser.add_argument("-url", "-u", type=str, help=tl("URL_ARGUMENT"), required=False) parser.add_argument( - "-list", "-l", type=str, help=tl("FILE_ARGUMENT"), required=False + "-u", "--url", type=str, help=tl("URL_ARGUMENT"), required=False + ) + parser.add_argument( + "-l", "--list", type=str, help=tl("FILE_ARGUMENT"), required=False ) parser.add_argument( - "-quality", "-q", - "-res", "-r", + "--quality", + "--resolution", type=str, help=tl("QUALITY_ARGUMENT"), required=False, ) parser.add_argument( - "-folder", "-f", type=str, help=tl("FOLDER_ARGUMENT"), required=False + "-f", "--folder", type=str, help=tl("FOLDER_ARGUMENT"), required=False ) parser.add_argument( - "-instance", "-i", type=str, help=tl("INSTANCE_ARGUMENT"), required=False + "-i", "--instance", type=str, help=tl("INSTANCE_ARGUMENT"), required=False ) parser.add_argument( - "-key", "-k", type=str, help=tl("APIKEY_ARGUMENT"), required=False + "-k", "--key", type=str, help=tl("APIKEY_ARGUMENT"), required=False ) parser.add_argument( - "-playlist", "-pl", + "--playlist", type=str, help=tl("PLAYLIST_ARGUMENT"), required=False, ) parser.add_argument( - "-filenameStyle", "-fs", + "--filenameStyle", type=str, help=tl("FILENAME_ARGUMENT"), required=False, choices=["classic", "pretty", "basic", "nerdy"], ) parser.add_argument( - "-audioFormat", "-af", + "--audioFormat", type=str, help=tl("AUDIO_ARGUMENT"), required=False, choices=["mp3", "ogg", "wav", "opus"], ) parser.add_argument( - "-youtubeVideoCodec", "-yvc", + "--youtubeVideoCodec", help=tl("YOUTUBE_VIDEO_CODEC_ARGUMENT"), required=False, choices=["vp9", "h264"], ) parser.add_argument( - "-show", "-s", + "--show", help=tl("SHOW_ARGUMENT"), action="store_true", ) parser.add_argument("-play", "-p", help=tl("PLAY_ARGUMENT"), action="store_true") parser.add_argument( - "-v", "-version", help=tl("VERSION_ARGUMENT"), action="store_true" + "-v", "--version", help=tl("VERSION_ARGUMENT"), action="store_true" ) parser.add_argument( - "-up", "-update", help=tl("UPDATE_ARGUMENT"), action="store_true" + "-up", "--update", help=tl("UPDATE_ARGUMENT"), action="store_true" ) args = parser.parse_args() - if args.v: + if args.version: try: print(tl("VERSION").format(version("pybalt"))) except Exception: print(tl("PACKAGE_NOT_FOUND")) return - if args.up: + if args.update: await check_updates() return if args.url_arg: @@ -101,20 +103,7 @@ async def _(): return api = Cobalt(api_instance=args.instance, api_key=args.key) if args.playlist: - await api.download( - url=args.playlist, - playlist=True, - path_folder=args.folder if args.folder else None, - quality=args.quality if args.quality else "1080", - filename_style=args.filenameStyle if args.filenameStyle else "pretty", - audio_format=args.audioFormat if args.audioFormat else "mp3", - youtube_video_codec=args.youtubeVideoCodec - if args.youtubeVideoCodec - else None, - show=args.show, - play=args.play, - ) - return + urls += [args.playlist] for url in urls: await api.download( url=url, diff --git a/pybalt/cobalt.py b/pybalt/cobalt.py index 6c1dd1d..81400c9 100644 --- a/pybalt/cobalt.py +++ b/pybalt/cobalt.py @@ -193,6 +193,9 @@ def __init__( self.skipped_instances = [] + def download_callback(self, *args, **kwargs): + """You can replace this function with your own callback function, called every ~2 seconds when downloading a file.""" + async def get_instance(self): """ Finds a good instance of Cobalt API and changes the API instance of this object to it. @@ -498,6 +501,7 @@ def shorten(s: str, additional_len: int = 0) -> str: async with ClientSession(headers=self.headers) as session: async with aopen(path.join(path_folder, filename), "wb") as f: try: + max_print_length, _ = 0, 0 progress_chars = ["⢎⡰", "⢎⡡", "⢎⡑", "⢎⠱", "⠎⡱", "⢊⡱", "⢌⡱", "⢆⡱"] progress_index = 0 total_size = 0 @@ -544,6 +548,13 @@ def shorten(s: str, additional_len: int = 0) -> str: f"\033[97m{info[:-2]}\033[94m{info[-2:]}\033[0m", end="", ) + self.download_callback( + total_size=total_size, + filename=filename, + path_folder=path_folder, + start_time=start_time, + speed_display=speed, + ) elapsed_time = time() - start_time info = f"[{round(total_size / 1024 / 1024, 2)}Mb \u2015 {round(elapsed_time, 2)}s] \u2713" print_line = shorten(result_path, additional_len=len(info)) diff --git a/setup.py b/setup.py index b5795d8..8fd1e85 100644 --- a/setup.py +++ b/setup.py @@ -57,4 +57,3 @@ def readme(): ], }, ) -