music/dl.py
2024-04-07 00:54:43 +03:00

133 lines
3.3 KiB
Python

import os, sys, json, subprocess, time;
from types import *;
from typing import *;
from dataclasses import dataclass
from argparse_dataclass import dataclass as argparser
from enum import IntEnum
@dataclass
class Process:
proc: subprocess.Popen[bytes]
name: str
url: str
file: str
genre: str
PROCESSES: List[Process] =[]
class Level(IntEnum):
ERROR = 0
WARN = 1
INFO = 2
DEBUG = 3
def level_to_prefix(l: Level) -> str:
match l:
case Level.ERROR:
return "\033[0;31m\033[1merror\033[0m"
case Level.WARN:
return "\033[1;33m\033[1mwarn\033[0m"
case Level.INFO:
return "\033[1;32m\033[1minfo\033[0m"
case Level.DEBUG:
return "\033[1;34m\033[1mdebug\033[0m"
class Logger:
level: Level = Level.INFO
def error(self, s: str):
if self.level >= Level.ERROR:
print(f"{level_to_prefix(Level.ERROR)}: {s}")
def warn(self, s: str):
if self.level >= Level.WARN:
print(f"{level_to_prefix(Level.WARN)}: {s}")
def info(self, s: str):
if self.level >= Level.INFO:
print(f"{level_to_prefix(Level.INFO)}: {s}")
def debug(self, s: str):
if self.level >= Level.DEBUG:
print(f"{level_to_prefix(Level.DEBUG)}: {s}")
LOGGER = Logger()
if os.name == 'nt':
YTDLP="yt-dlp.exe"
else:
YTDLP="yt-dlp"
@dataclass
class ManifestSong:
name: str
url: str
type ManifestGenre = List[ManifestSong]
type Manifest = Dict[str, ManifestGenre]
@argparser
class CliArgs:
manifest: str = "./manifest.json"
output_dir: str = "./"
debug: bool = False
def main(cliargs: CliArgs) -> int:
if cliargs.debug:
LOGGER.level = Level.DEBUG
manifest: Manifest = get_manifest(cliargs)
count = 0
for (genre, v) in manifest.items():
for song in v:
if len(PROCESSES) > 10:
wait_for_procs(10)
if download_song(genre, song["name"], song["url"]):
count += 1
wait_for_procs()
LOGGER.info(f"Downloaded {count} new songs")
def download_song(dir: str, name: str, url: str) -> bool:
if not os.path.isdir(dir):
os.mkdir(dir)
outfile=f"{dir}/{name}.m4a"
if os.path.isfile(outfile):
LOGGER.debug(f"Already downloaded {outfile} ({url}), skipping")
return False
LOGGER.info(f"Downloading {outfile} ({url})")
proc = subprocess.Popen(
[YTDLP, "-x", "--audio-format", "m4a","-o", outfile, url],
stdout=open(os.devnull, 'wb')
)
PROCESSES.append(Process(proc, name, url, outfile, dir))
return True
def wait_for_procs(until: int = 0):
while len(PROCESSES) > until:
for proc in PROCESSES:
retcode = proc.proc.poll()
if retcode is not None: # Process finished.
PROCESSES.remove(proc)
LOGGER.info(f"Finished downloading {proc.file}")
break
else: # No process is done, wait a bit and check again.
time.sleep(.1)
continue
def get_manifest(cliargs: CliArgs) -> Manifest:
with open(cliargs.manifest, "r") as f:
data = f.read()
return json.loads(data)
if __name__ == "__main__":
ret = main(CliArgs.parse_args(sys.argv[1:]))
sys.exit(ret)