beets/beetsplug/bench.py

"""Some simple performance benchmarks for beets."""

import cProfile
import timeit

from beets import importer, plugins, ui
from beets.autotag import tag_album
from beets.plugins import BeetsPlugin
from beets.util.functemplate import Template
from beets.util.pathformats import PF_KEY_DEFAULT
from beetsplug._utils import vfs


def aunique_benchmark(lib, prof):
    def _build_tree():
        vfs.libtree(lib)

    # Measure path generation performance with %aunique{} included.
    lib.path_formats = [
        (
            PF_KEY_DEFAULT,
            Template("$albumartist/$album%aunique{}/$track $title"),
        )
    ]
    if prof:
        cProfile.runctx(
            "_build_tree()",
            {},
            {"_build_tree": _build_tree},
            "paths.withaunique.prof",
        )
    else:
        interval = timeit.timeit(_build_tree, number=1)
        print("With %aunique:", interval)

    # And with %aunique replaced with a "cheap" no-op function.
    lib.path_formats = [
        (PF_KEY_DEFAULT, Template("$albumartist/$album%lower{}/$track $title"))
    ]
    if prof:
        cProfile.runctx(
            "_build_tree()",
            {},
            {"_build_tree": _build_tree},
            "paths.withoutaunique.prof",
        )
    else:
        interval = timeit.timeit(_build_tree, number=1)
        print("Without %aunique:", interval)


def match_benchmark(lib, prof, query=None, album_id=None):
    # If no album ID is provided, we'll match against a suitably huge
    # album.
    if not album_id:
        album_id = "9c5c043e-bc69-4edb-81a4-1aaf9c81e6dc"

    # Get an album from the library to use as the source for the match.
    items = lib.albums(query).get().items()

    # Ensure fingerprinting is invoked (if enabled).
    plugins.send(
        "import_task_start",
        task=importer.ImportTask(None, None, items),
        session=importer.ImportSession(lib, None, None, None),
    )

    # Run the match.
    def _run_match():
        tag_album(items, search_ids=[album_id])

    if prof:
        cProfile.runctx(
            "_run_match()", {}, {"_run_match": _run_match}, "match.prof"
        )
    else:
        interval = timeit.timeit(_run_match, number=1)
        print("match duration:", interval)


class BenchmarkPlugin(BeetsPlugin):
    """A plugin for performing some simple performance benchmarks."""

    def commands(self):
        aunique_bench_cmd = ui.Subcommand(
            "bench_aunique", help="benchmark for %aunique{}"
        )
        aunique_bench_cmd.parser.add_option(
            "-p",
            "--profile",
            action="store_true",
            default=False,
            help="performance profiling",
        )
        aunique_bench_cmd.func = lambda lib, opts, args: aunique_benchmark(
            lib, opts.profile
        )

        match_bench_cmd = ui.Subcommand(
            "bench_match", help="benchmark for track matching"
        )
        match_bench_cmd.parser.add_option(
            "-p",
            "--profile",
            action="store_true",
            default=False,
            help="performance profiling",
        )
        match_bench_cmd.parser.add_option(
            "-i", "--id", default=None, help="album ID to match against"
        )
        match_bench_cmd.func = lambda lib, opts, args: match_benchmark(
            lib, opts.profile, args, opts.id
        )

        return [aunique_bench_cmd, match_bench_cmd]