add make_mines_fake.py & test data

make_mines_fake.py

@@ -0,0 +1,508 @@
+R"""
+Add a short fake region on every isolated mine
+to prevent them from being hit during gameplay.
+
+This script only works on SSC files.
+
+Usage examples:
+
+    # Preview changes & detect any errors
+    python make_mines_fake.py "C:\StepMania\Songs\My Pack\My Song" --dry-run
+
+    # Preview changes with errors & warnings suppressed
+    python make_mines_fake.py "C:\StepMania\Songs\My Pack\My Song" --dry-run --allow-simultaneous --allow-split-timing --ignore-sm
+
+    # Commit changes with only warnings suppressed
+    python make_mines_fake.py "C:\StepMania\Songs\My Pack\My Song" --ignore-sm
+"""
+import argparse
+import bisect
+from dataclasses import dataclass, field
+from decimal import Decimal
+import heapq
+import os
+import sys
+from typing import Iterator, List, NamedTuple, Optional, Union
+
+import simfile
+import simfile.dir
+from simfile.notes import NoteData, NoteType
+from simfile.notes.group import group_notes, SameBeatNotes
+from simfile.timing import Beat, BeatValue, BeatValues
+
+
+####################
+# Script arguments #
+####################
+
+
+class MakeMinesFakeArgs:
+    """Stores the command-line arguments for this script."""
+
+    simfile: str
+    dry_run: bool
+    allow_split_timing: bool
+    allow_simultaneous: bool
+    ignore_sm: bool
+
+
+def argparser():
+    """Get an ArgumentParser instance for this command-line script."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("simfile", type=str, help="path to the simfile to modify")
+    parser.add_argument(
+        "-d",
+        "--dry-run",
+        action=argparse.BooleanOptionalAction,
+        help="preview changes without writing the file",
+    )
+    parser.add_argument(
+        "--allow-split-timing",
+        action=argparse.BooleanOptionalAction,
+        help="if necessary, create split timing for charts to avoid interfering with other charts",
+    )
+    parser.add_argument(
+        "--allow-simultaneous",
+        action=argparse.BooleanOptionalAction,
+        help="leave mines that occur on the same beat as notes alone",
+    )
+    parser.add_argument(
+        "--ignore-sm",
+        action=argparse.BooleanOptionalAction,
+        help="do not warn when an SM file is present alongside the SSC file",
+    )
+    return parser
+
+
+#####################
+# Utility functions #
+#####################
+
+
+def whichchart(
+    ssc: simfile.ssc.SSCSimfile, chart: Union[int, simfile.ssc.SSCChart]
+) -> str:
+    """
+    Identify the chart by its difficulty,
+    as well as its stepstype if multiple are present in the simfile.
+    Intended for human-readable output.
+    """
+    if isinstance(chart, int):
+        chartobj_ = ssc.charts[chart]
+        chartobj: simfile.ssc.SSCChart = chartobj_  # typing workaround
+    else:
+        chartobj = chart
+
+    stepstypes = set(ch.stepstype for ch in ssc.charts)
+    if len(stepstypes) > 1:
+        return f"{chartobj.stepstype} {chartobj.difficulty}"
+    else:
+        return chartobj.difficulty or ""
+
+
+def whichtarget(
+    ssc: simfile.ssc.SSCSimfile,
+    target: Union[simfile.ssc.SSCSimfile, simfile.ssc.SSCChart],
+) -> str:
+    """
+    Identify the target as a chart (see :func:`whichchart`) or simfile.
+    Intended for human-readable output.
+
+    Returns a string starting with "the ".
+    """
+    if isinstance(target, simfile.ssc.SSCChart):
+        return f"the {whichchart(ssc, target)} chart"
+    else:
+        return "the simfile"
+
+
+def hastiming(chart: simfile.ssc.SSCChart) -> bool:
+    """
+    Detect whether the chart has its own timing data.
+    """
+    return chart.bpms is not None
+
+
+def splittiming(ssc: simfile.ssc.SSCSimfile, chart: simfile.ssc.SSCChart) -> None:
+    """
+    Copy timing data from the SSC simfile to its chart.
+    """
+    for td in (
+        "STOPS",
+        "DELAYS",
+        "BPMS",
+        "WARPS",
+        "LABELS",
+        "TIMESIGNATURES",
+        "TICKCOUNTS",
+        "COMBOS",
+        "SPEEDS",
+        "SCROLLS",
+        "FAKES",
+    ):
+        chart[td] = ssc[td]
+
+
+################
+# Script logic #
+################
+
+
+class BeatChartIndex(NamedTuple):
+    beat: Beat
+    chart_index: int
+
+
+class SameBeatMineAndNote(NamedTuple):
+    beat: Beat
+    mine_chart_index: int
+    note_chart_index: int
+
+
+@dataclass
+class NotesAndMines:
+    """
+    Note & mine data across all charts of a simfile.
+    """
+
+    note_positions: List[BeatChartIndex] = field(default_factory=list)
+    """
+    All of the non-fake, non-mine note object positions in each chart.
+    Sorted first by beat, then by chart index.
+    """
+
+    mine_positions: List[BeatChartIndex] = field(default_factory=list)
+    """
+    All of the mine note object positions in each chart.
+    Sorted first by beat, then chart index.
+    """
+
+    must_allow_simultaneous: List[SameBeatMineAndNote] = field(default_factory=list)
+    """All of the same-beat mine & note pairs within the same chart."""
+
+    must_allow_split_timing: List[SameBeatMineAndNote] = field(default_factory=list)
+    """All of the same-beat mine & note pairs across different charts."""
+
+
+class SameBeatMineAndNoteError(Exception):
+    """
+    Raised if a mine & note occur on the same beat
+    and such an occurrence isn't explicitly allowed by the input arguments.
+    Stringifies to a multi-line, human-readable error message.
+    """
+
+    def __init__(
+        self,
+        ssc: simfile.ssc.SSCSimfile,
+        must_allow_simultaneous: List[SameBeatMineAndNote],
+        must_allow_split_timing: List[SameBeatMineAndNote],
+    ):
+        self.ssc = ssc
+        self.must_allow_simultaneous = must_allow_simultaneous
+        self.must_allow_split_timing = must_allow_split_timing
+
+    def _stringify_simultaneous(self, s: SameBeatMineAndNote):
+        if s.mine_chart_index == s.note_chart_index:
+            return f"b{s.beat} in {whichchart(self.ssc, s.mine_chart_index)}"
+        else:
+            return f"b{s.beat} in {whichchart(self.ssc, s.mine_chart_index)} and {whichchart(self.ssc, s.note_chart_index)}"
+
+    def _stringify_simultaneous_list(self, sl: List[SameBeatMineAndNote]):
+        return "".join(f"    {self._stringify_simultaneous(s)}\n" for s in sl)
+
+    def __str__(self):
+        return (
+            f"ERROR: There are simultaneous mines and notes in {self.ssc.title};\n"
+            "you will have to either update the chart or pass additional flags to the script\n"
+            "in order to remedy this error.\n"
+            + (
+                "\n"
+                "Simultaneous mine & note in the same chart (ignore with --allow-simultaneous):\n"
+                f"{self._stringify_simultaneous_list(self.must_allow_simultaneous)}"
+                "Ignoring these occurrences will leave the mines on these beats hittable,\n"
+                "which may surprise players. Consider relocating these mines instead.\n"
+                if self.must_allow_simultaneous
+                else ""
+            )
+            + (
+                "\n"
+                "Simultaneous mine & note in different charts (fix with --allow-split-timing):\n"
+                f"{self._stringify_simultaneous_list(self.must_allow_split_timing)}"
+                "Note that split timing makes it easy to mess up the timing data if you are\n"
+                "still making changes to the chart. Use this feature with caution.\n"
+                if self.must_allow_split_timing
+                else ""
+            )
+        )
+
+
+def find_chart_with_existing_item(
+    positions: List[BeatChartIndex],
+    position: BeatChartIndex,
+) -> Optional[int]:
+    """
+    If `positions` has an item on the same beat as `position`,
+    returns the chart index of the found item.
+    """
+    if not positions:
+        return None
+    index = bisect.bisect_left(positions, position.beat, key=lambda pos: pos.beat)
+    if index == len(positions):
+        return None
+    item = positions[index]
+    if item.beat == position.beat:
+        return item.chart_index
+
+
+def append_same_beat_items(
+    nm: NotesAndMines,
+    chart_positions: List[BeatChartIndex],
+    position: BeatChartIndex,
+    *,
+    position_is_mine: bool,
+) -> None:
+    """
+    If the given position occurs in the other type of items (between notes and mines),
+    append a SameBeatMineAndNote to the appropriate NotesAndMines `must_allow` list.
+    """
+    items_from_other_charts = (
+        nm.note_positions if position_is_mine else nm.mine_positions
+    )
+
+    # Look for a note on the same beat in a previously indexed chart
+    if (
+        other_chart := find_chart_with_existing_item(items_from_other_charts, position)
+    ) is not None:
+        mine_chart_index = position.chart_index if position_is_mine else other_chart
+        note_chart_index = other_chart if position_is_mine else position.chart_index
+        nm.must_allow_split_timing.append(
+            SameBeatMineAndNote(
+                beat=position.beat,
+                mine_chart_index=mine_chart_index,
+                note_chart_index=note_chart_index,
+            )
+        )
+
+    # Look for a note on the same beat in this chart
+    if find_chart_with_existing_item(chart_positions, position) is not None:
+        nm.must_allow_simultaneous.append(
+            SameBeatMineAndNote(
+                beat=position.beat,
+                mine_chart_index=position.chart_index,
+                note_chart_index=position.chart_index,
+            )
+        )
+
+
+def get_notes_and_mines(ssc: simfile.ssc.SSCSimfile) -> NotesAndMines:
+    """
+    Populate & return a NotesAndMines instance for the simfile.
+    """
+    output = NotesAndMines()
+
+    for c, chart in enumerate(ssc.charts):
+        chart_note_positions: List[BeatChartIndex] = []
+        """All of the non-fake, non-mine note object positions in this chart. Sorted by beat."""
+
+        chart_mine_positions: List[BeatChartIndex] = []
+        """All of the mine note object positions in this chart. Sorted by beat."""
+
+        nd = NoteData(chart)
+        for note in nd:
+            position = BeatChartIndex(beat=note.beat, chart_index=c)
+            match note.note_type:
+                case NoteType.MINE:
+                    append_same_beat_items(
+                        output, chart_note_positions, position, position_is_mine=True
+                    )
+                    chart_mine_positions.append(position)
+
+                case NoteType.FAKE:
+                    pass
+
+                case _:  # Any hittable or scoreable "note"
+                    append_same_beat_items(
+                        output, chart_mine_positions, position, position_is_mine=False
+                    )
+                    chart_note_positions.append(position)
+
+        output.note_positions = list(
+            heapq.merge(output.note_positions, chart_note_positions)
+        )
+        output.mine_positions = list(
+            heapq.merge(output.mine_positions, chart_mine_positions)
+        )
+
+    return output
+
+
+def maybe_raise_simultaneous_error(
+    ssc: simfile.ssc.SSCSimfile,
+    args: MakeMinesFakeArgs,
+    nm: NotesAndMines,
+) -> None:
+    """
+    Raise :class:`SameBeatMineAndNoteError`
+    only if the script arguments don't allow for same-beat mine & note pairs
+    that were found in the simfile.
+    """
+    simultaneous = [] if args.allow_simultaneous else nm.must_allow_simultaneous
+    split_timing = [] if args.allow_split_timing else nm.must_allow_split_timing
+
+    if simultaneous or split_timing:
+        raise SameBeatMineAndNoteError(
+            ssc,
+            must_allow_simultaneous=simultaneous,
+            must_allow_split_timing=split_timing,
+        )
+
+
+def charts_with_mines(
+    ssc: simfile.ssc.SSCSimfile, nm: NotesAndMines
+) -> Iterator[simfile.ssc.SSCChart]:
+    """
+    Yield only the charts that have mines.
+    Sorted by chart index.
+    """
+    chart_indexes = set(mine_pos.chart_index for mine_pos in nm.mine_positions)
+    for chart_index in sorted(chart_indexes):
+        yield ssc.charts[chart_index]
+
+
+def make_mines_fake(ssc: simfile.ssc.SSCSimfile, args: MakeMinesFakeArgs) -> List[str]:
+    """
+    Add a short fake segment on each mine in each chart in the simfile,
+    updating the in-memory simfile object.
+
+    Returns a list of actions (human-readable strings)
+    that were taken on the simfile.
+
+    Raises :class:`SameBeatMineAndNoteError`
+    if a mine & note are found on the same beat
+    and the arguments don't explicitly allow it.
+    """
+    actions = []
+    nm = get_notes_and_mines(ssc)
+    maybe_raise_simultaneous_error(ssc, args, nm)
+
+    for chart in charts_with_mines(ssc, nm):
+        if hastiming(chart):
+            fakes_target = chart
+        elif len(nm.must_allow_split_timing) > 0:
+            actions.append(
+                f"copy timing data from the simfile to the {whichchart(ssc, chart)} chart"
+            )
+            splittiming(ssc, chart)
+            fakes_target = chart
+        else:
+            fakes_target = ssc
+
+        fakes = (
+            BeatValues.from_str(fakes_target.fakes)
+            if fakes_target.fakes
+            else BeatValues()
+        )
+
+        for note_group in group_notes(
+            NoteData(chart),
+            same_beat_notes=SameBeatNotes.JOIN_ALL,
+        ):
+            if any(note.note_type == NoteType.MINE for note in note_group):
+                beat = note_group[0].beat
+
+                if not all(note.note_type == NoteType.MINE for note in note_group):
+                    # This case should've been caught already
+                    assert args.allow_simultaneous
+                    continue
+
+                already_fake = False
+                for fake_ in fakes:
+                    fake: BeatValue = fake_  # typing workaround
+
+                    # Convert the decimal value to string first
+                    # to force Beat() to quantize it to a 192nd
+                    fake_end_beat = fake.beat + Beat(str(fake.value))
+
+                    # Skip mines that are already inside a fake segment
+                    if fake.beat <= beat < fake_end_beat:
+                        already_fake = True
+                        break
+
+                    # Stop checking fake segments past the current mine's beat
+                    # Assumption: fake region start & end beats are strictly increasing
+                    if fake.beat > beat:
+                        already_fake = False
+                        break
+
+                if already_fake:
+                    continue
+
+                actions.append(
+                    f"add a short fake region on b{beat} to {whichtarget(ssc, fakes_target)}"
+                )
+                bisect.insort(
+                    fakes,
+                    BeatValue(beat=beat, value=Decimal(str(Beat.tick()))),
+                    key=lambda bv: bv.beat,
+                )
+
+        fakes_target.fakes = str(fakes)
+
+    return actions
+
+
+def main(argv) -> int:
+    """
+    Run the script & return an exit code (0 for success, nonzero for error).
+    """
+    error = False
+
+    # Parse command-line arguments
+    args = argparser().parse_args(argv[1:], namespace=MakeMinesFakeArgs())
+
+    if os.path.isdir(args.simfile):
+        sd = simfile.dir.SimfileDirectory(args.simfile)
+        if sd.ssc_path:
+            args.simfile = sd.ssc_path
+        else:
+            raise ValueError("simfile directory has no SSC file")
+
+        if sd.sm_path and not args.ignore_sm:
+            print(
+                "WARNING: there is an SM file in this directory that won't be touched.\n"
+                "Either delete the SM file or pass --ignore-sm to suppress this warning.\n"
+            )
+
+    input_filename = args.simfile
+    backup_filename = input_filename + "~"
+
+    with simfile.mutate(input_filename, backup_filename=backup_filename) as sf:
+        if isinstance(sf, simfile.ssc.SSCSimfile):
+            try:
+                actions = make_mines_fake(sf, args)
+            except SameBeatMineAndNoteError as e:
+                actions = []
+                error = True
+                print(str(e))
+
+            if actions:
+                print("Actions taken:\n" + "".join(f"    {a}\n" for a in actions))
+            else:
+                print("No actions taken")
+                raise simfile.CancelMutation
+        else:
+            raise TypeError("fakes require an SSC file")
+
+        if args.dry_run:
+            print("Not writing changes for dry run")
+            raise simfile.CancelMutation
+        else:
+            print(
+                f"Writing changes to {input_filename} & backing up original file to {backup_filename}"
+            )
+
+    return 1 if error else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv))