main.py

#!/usr/bin/env python
# PYTHON_ARGCOMPLETE_OK

import sys
import libopensonic  # pyright: ignore[reportMissingTypeStubs]
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
import argparse
from netrc import netrc
import re
from sys import stderr
import logging
from libopensonic.media.media_types import Child
from tqdm import tqdm
from magic import Magic
from glob import glob
import argcomplete
import importlib.metadata
from importlib.metadata import PackageNotFoundError
import pathlib
import mimetypes

# 🟰 definitions

mime = Magic(mime=True)

if os.name == "nt":

    # mercy for w*ndows users
    def sanitize(string: str) -> str:
        """
        Edit text to be suitable for use in a windows path. Replaces invalid characters with `-`.

        :param string: input text
        :type string: str
        :return: edited text
        :rtype: str
        """
        return re.sub(r"/|\\|<|>|:|\"|\||\?|\*|\x00", "-", string)

else:

    def sanitize(string: str) -> str:
        """
        Edit text to be suitable for use in a *nix path. Replaces invalid characters with `-`.

        :param string: input text
        :type string: str
        :return: edited text
        :rtype: str
        """
        return re.sub(r"/|\\|:", "-", string)


def path(
    artist: str | None,
    album: str | None = None,
    song: str | None = None,
    extension: str | None = None,
) -> str:
    """
    Generate a `sanitize`d path for an artist directory, album directory, or song.

    For example, "5/3" from the album *Becoming an Entity*, by The Doozers:
    - `path("The Doozers")` → `"The Doozers"`
    - `path("The Doozers", "Becoming an Entity")` → `"The Doozers/Becoming an Entity"`
    - `path("The Doozers", "Becoming an Entity", "5/3", "mp3")` → `"The Doozers/Becoming an Entity/5-3.mp3"`
    - `path(None, None, "5/3", "mp3")` → `"5-3.mp3"`

    :param artist: the name of an artist
    :type artist: str | None
    :param album: the name of said artist's album
    :type album: str | None
    :param song: the title of said album's song
    :type song: str | None
    :param extension: the desired file extension for the song
    :type song: str | None
    :return: a `sanitize`d path
    :rtype: str
    """

    if artist == None:
        assert album == None and not song == None and not extension == None

        return sanitize(song) + "." + extension

    result = sanitize(artist)

    if album:
        result += "/" + sanitize(album)

        if song:
            assert extension

            result += "/" + sanitize(song) + extension

    return result


# ⚙️ configure arguments

more_info_message = "ℹ️ \0 See https://github.com/solanto/downsonic for more information."

argument_parser = argparse.ArgumentParser(
    description="🎷 downsonic is a little command-line utility made to help download all of a user's music from an OpenSubsonic server.",
    epilog=more_info_message
)

argument_parser.add_argument(
    "source",
    help="OpenSubsonic server to download music from; [http[s]://]HOST[:PORT], where port defaults to 80 for http, 443 for https, and 8080 when unspecified",
)
argument_parser.add_argument("destination", help="destination directory")
argument_parser.add_argument(
    "--netrc-file",
    help="path to a netrc file with login credentials; defaults to `~/.netrc`",
)
argument_parser.add_argument(
    "-u", "--user", help="username for server login ⚠️ \0 see help on `--password`"
)
argument_parser.add_argument(
    "-p",
    "--password",
    help="password for server login ⚠️ \0 avoid supplying your passwords through the terminal in plaintext, and instead consider using a netrc file",
)
argument_parser.add_argument(
    "-b",
    "--bitrate",
    type=int,
    help="target bitrate for transcoded files, in kbps; unspecified or `0` set no limit",
)
argument_parser.add_argument(
    "-F",
    "--format",
    help="audio format; navidrome servers, for example, support `mp3`, `flac`, `aac`, and `raw` (no transcoding)",
)
argument_parser.add_argument(
    "-e",
    "--extension",
    help="override the extension of all audio files; when unspecified, each file's extension will be inferred automatically",
)
argument_parser.add_argument(
    "-t",
    "--threads",
    help="maximum number of threads (and parallel network connections) to use while downloading",
    default=None,
)
argument_parser.add_argument(
    "-f",
    "--force",
    help="(re)download songs even if they're already in the destination directory",
    action="store_true",
)
argument_parser.add_argument(
    "-v",
    "--verbosity",
    action="count",
    default=0,
    help="how much logging to show; `-v` for critical errors (🛑), `-vv` for recoverable errors (⛔️), `-vvv` for warnings (⚠️ ), `-vvvv` for info (default), and `-vvvvv` for debug (🪲)",
)
argument_parser.add_argument(
    "--non-interactive",
    action="store_true",
    help="don't show dynamic UI elements, like progress bars",
)
argument_parser.add_mutually_exclusive_group().add_argument(
    "--completions",
    action="store_true",
    help="print instructions to enable shell autocompletion and exit",
)
argument_parser.add_mutually_exclusive_group().add_argument(
    "-V",
    "--version",
    action="store_true",
    help="print the program's version and exit",
)

executable = sys.argv[0].split(os.path.sep)[-1]

argument_parser.usage = f"""{argument_parser.format_usage().split(": ")[1]}
examples:
  {executable} music.server.local ~/Music
  {executable} https://music.server.me ~/Music --netrc-file ~/.another-netrc
  {executable} https://music.server.me:1234 ~/Music -F mp3 -b 320"""

argcomplete.autocomplete(argument_parser)

# 🪵 configure logging


class LevelFormatter(logging.Formatter):
    """
    A custom `logging.Formatter` extension that allows for individual formatting per log level.
    """

    def __init__(self, formats):
        super().__init__()
        self.formats = formats

    def format(self, record):
        # Get the appropriate format string for the current log level
        format_string = self.formats[record.levelno]

        # If a specific format is defined for the level, use it
        if format_string:
            self._style._fmt = format_string
        else:
            # Fallback to a default format if no specific format is found for the level
            self._style._fmt = self.formats["default"]

        return super().format(record)


streamHandler = logging.StreamHandler()

streamHandler.setFormatter(
    LevelFormatter(
        {
            logging.DEBUG: "🪲  %(module)s:%(lineno)d · %(message)s",
            logging.INFO: "%(message)s",
            logging.WARNING: "⚠️  %(message)s",
            logging.ERROR: "⛔️  %(message)s",
            logging.CRITICAL: "🛑 %(message)s",
            "default": "%(message)s",
        }
    )
)

logging.root.addHandler(streamHandler)


def run():
    """
    Main functionality.
    """

    meter: tqdm | None = None

    # 📖 parse & validate arguments

    if "--version" in sys.argv[1:] or "-V" in sys.argv[1:]:
        try:
            print(f"🎻 downsonic {importlib.metadata.version("downsonic")}")
        except PackageNotFoundError:
            print(
                f"🤓 downsonic is not installed as a package; check {
                    pathlib.Path(__file__).parent.resolve().parent.parent.joinpath("pyproject.toml")
                }:project.version"
            )
        
        print("\n" + more_info_message)

        exit(0)


    if "--completions" in sys.argv[1:]:
        print( # fmt: skip
"""1️⃣  First, install argcomplete (https://kislyuk.github.io/argcomplete/).

On pipx systems:

    pipx install argcomplete

On pip systems:

    pip install argcomplete

2️⃣  Then, activate argcomplete.

    activate-global-python-argcomplete
    
3️⃣  Finally, register completions for downsonic.

On bash and zsh systems:

    eval "$(register-python-argcomplete downsonic)"

On fish systems:

    register-python-argcomplete --shell fish downsonic > ~/.config/fish/completions/downsonic.fish

For other shells, see https://kislyuk.github.io/argcomplete/#support-for-other-shells.
"""
        )
        exit(0)

    arguments = argument_parser.parse_args()

    logging.root.setLevel(
        max(6 - arguments.verbosity, 1) * 10 if arguments.verbosity else 20,
    )

    if arguments.threads and not arguments.threads.isdigit():
        logging.critical("invalid number of threads")
        exit(1)

    try:
        authenticator = netrc(arguments.netrc_file).authenticators(arguments.source)

        if not authenticator:
            raise KeyError()

        user, _, password = authenticator
    except (FileNotFoundError, KeyError):
        if not arguments.user or not arguments.password:
            logging.critical(
                "unable to get user and password from netrc file or arguments"
            )
            exit(1)

        user = arguments.user
        password = arguments.password

    url = re.search(
        r"^(?:(https?):\/\/)?((?!.+:\/{0,2}$).+?)(?::(\d+))?$", arguments.source
    )

    if not url or not url[2]:
        logging.critical("unable to parse url")
        exit(1)

    protocol, host, port_string = url.groups()

    # parse explicit port or infer from protocol
    port = (
        int(port_string)
        if port_string
        else 443 if protocol == "https" else 80 if protocol == "http" else 8080
    )

    # parse explicit procotol or infer from port
    protocol = protocol if protocol else "https" if port == 443 else "http"

    threads = int(arguments.threads) if arguments.threads else None

    if not os.path.isdir(arguments.destination):
        logging.critical("destination folder does not exist")
        exit(1)

    # ⚓️ initialize opensubsonic client & helpers

    client = libopensonic.Connection(protocol + "://" + host, user, password, port)

    def write_song(song: Child):
        """
        From OpenSubsonic data about a song, download the song to its appropriate directory.

        :param song: OpenSubsonic song data
        :type song: Child
        """

        if meter and logging.root.level <= logging.INFO:
            meter.write(song.title, stderr)

        response = client.stream(song.id, arguments.bitrate, arguments.format)

        if arguments.extension:
            extension = "." + arguments.extension
        else:  # infer extension from mimetype
            mimetype = response.headers.get("content-type") or mime.from_buffer(response.content)

            extension = (
                     ".mp3"  if mimetype == "audio/mpeg"
                else ".m4a"  if re.match(r"^audio\/((.*\W)?(aac|mp4a?)(\W.*)?)$", mimetype)
                else ".flac" if re.match(r"^audio\/(x-)?flac$", mimetype)
                else ".wav"  if mimetype == "audio/wav"
                else ".ogg"  if mimetype == "audio/ogg"
                else ".aiff" if re.match(r"^audio\/(x-)?aiff$", mimetype)
                else ".wma"  if mimetype == "video/x-ms-asf" or mimetype == "audio/x-ms-wma"
                else guess   if (guess := mimetypes.guess_extension(mimetype))
                else ""
            )

        song_path = (
            path(song.artist, song.album, song.title, extension)
            if song.artist
            else path(None, None, song.title, extension)
        )

        with open(song_path, "wb") as file:
            file.write(response.content)

    # 📂 work in destination directory

    os.chdir(arguments.destination)

    # 🗄️ prepare subdirectories

    indexes = client.get_indexes().index

    if not indexes:
        logging.warning("server returned no indexes")
        exit(0)

    all_songs = []

    for index in indexes:
        if index.artist:
            for artist in index.artist:
                os.makedirs(path(artist.name), exist_ok=True)

                albums: list[Child] | None = client.get_music_directory(artist.id).child

                if albums:
                    for album in albums:
                        os.makedirs(path(artist.name, album.title), exist_ok=True)

                        songs = client.get_music_directory(album.id).child

                        if songs:
                            for song in songs:
                                song_path = (
                                    path(song.artist, song.album, song.title, "*")
                                    if song.artist
                                    else path(None, None, song.title, "*")
                                )

                                if arguments.force or not glob(song_path):
                                    all_songs.append(song)

    # ⬇️ download songs

    with ThreadPoolExecutor(threads) as executor:
        futures = [executor.submit(write_song, song) for song in all_songs]

        if logging.root.level <= logging.INFO and not arguments.non_interactive:
            meter = tqdm(as_completed(futures), total=len(futures))
            meter.write("🎺 downloading songs:", stderr)

            results = []

            for future in meter:
                results.append(future.result())


def main():
    """A wrapper for `run` that catches `KeyboardInterrupt` exceptions."""

    try:
        run()
    except KeyboardInterrupt:
        logging.critical(
            "interrupted; you may need to interrupt again to stop all threads"
        )


if __name__ == "__main__":
    main()