simfile-scripts/sort_by_difficulty.py

R"""
Change the title of every simfile in a pack
so that they are sorted by difficulty in StepMania.

This script finds the hardest chart of a given stepstype (dance-single by default)
and puts its meter (difficulty number) between brackets at the start of the title.

Usage examples:

    # Sort a pack by difficulty
    python sort_by_difficulty.py "C:\StepMania\Songs\My Pack"

    # Unsort by difficulty (remove the title prefixes)
    python sort_by_difficulty.py -r "C:\StepMania\Songs\My Pack"

    # Customize stepstype and digits
    python sort_by_difficulty.py -s dance-double -d 3 "C:\StepMania\My Pack"
"""
import argparse
import sys
from typing import Optional, Sequence

import simfile
import simfile.dir


class SortByDifficultyArgs:
    """Stores the command-line arguments for this script."""

    pack: str
    stepstype: str
    digits: int
    remove: bool


def argparser():
    """Get an ArgumentParser instance for this command-line script."""
    parser = argparse.ArgumentParser()
    parser.add_argument("pack", type=str, help="path to the pack to modify")
    parser.add_argument("-s", "--stepstype", type=str, default="dance-single")
    parser.add_argument(
        "-d",
        "--digits",
        type=int,
        default=2,
        help="minimum digits (will add leading zeroes)",
    )
    parser.add_argument(
        "-r",
        "--remove",
        action=argparse.BooleanOptionalAction,
        help="remove meter prefix",
    )
    return parser


def hardest_chart(
    charts: Sequence[simfile.types.Chart], stepstype: str
) -> Optional[simfile.types.Chart]:
    """
    Find & return the hardest chart (numerically) of a given stepstype.

    Returns None if there are no charts matching the stepstype.
    """
    return max(
        [c for c in charts if c.stepstype == stepstype],
        key=lambda c: int(c.meter or "1"),
        default=None,
    )


def prefix_title_with_meter(simfile_path: str, args: SortByDifficultyArgs):
    """
    Add (or remove) a numeric prefix to the simfile's title.

    This saves the updated simfile to its original location
    and writes a backup copy with a ~ appended to the filename.
    """
    # You could specify output_filename here to write the updated file elsewhere
    with simfile.mutate(
        input_filename=f"{simfile_path}",
        backup_filename=f"{simfile_path}~",
    ) as sf:
        print(f"Processing {simfile_path}")

        # It's very unlikely for the title property to be blank or missing.
        # This is mostly to satisfy type-checkers.
        current_title = sf.title or ""

        if args.remove:
            # Look for a number in brackets at the start of the title
            if current_title.startswith("["):
                open_bracket_index = current_title.find("[")
                close_bracket_index = current_title.find("]")
                bracketed_text = current_title[
                    open_bracket_index + 1 : close_bracket_index
                ]
                if bracketed_text.isnumeric():
                    # Remove the bracketed number from the title
                    sf.title = current_title[close_bracket_index + 1 :].lstrip(" ")
        else:
            # Find the hardest chart (numerically) within a stepstype
            # and use it to prefix the title
            chart = hardest_chart(sf.charts, args.stepstype)

            # Skip this simfile if there were no charts for the stepstype.
            # Nothing will be written to disk in this case.
            if not chart:
                raise simfile.CancelMutation

            # It's very unlikely for the meter property to be blank or missing.
            # This is mostly to satisfy type-checkers.
            meter = chart.meter or "1"

            # Put the meter at the start of the title,
            # filling in leading zeros per arguments
            sf.title = f"[{meter.zfill(args.digits)}] {current_title}"


def main(argv):
    # Parse command-line arguments
    args = argparser().parse_args(argv[1:], namespace=SortByDifficultyArgs())

    # Iterate over SimfileDirectory objects from the pack
    # so that we can easily get the .sm and/or .ssc paths
    for simfile_dir in simfile.dir.SimfilePack(args.pack).simfile_dirs():

        # Try to update whichever formats exist
        for simfile_path in [simfile_dir.sm_path, simfile_dir.ssc_path]:
            if simfile_path:
                prefix_title_with_meter(simfile_path, args)


if __name__ == "__main__":
    main(sys.argv)