harness.py

Part of Balatro Feel.

  1"""Scripted-input harness: exercises the menu, hover, select, drag, and play.
  2
  3Drives the real input pipeline through ``InputSimulator`` (the same state writes and
  4event propagation the platform adapters perform) and saves one screenshot per stage:
  5
  6    01_menu.png      landing screen, before Play is pressed
  7    02_idle.png      rest pose, full hand fanned out
  8    03_hover.png     hovering the middle card: scale-up + parallax tilt
  9    04_select.png    that card clicked: lifts, punches, rises above its neighbours
 10    05_drag.png      leftmost card dragged right, slots swapping as it crosses
 11    06_dropped.png   dropped into its new slot, hand re-fanned
 12    07_play.png      two cards selected and played: staged pulse
 13
 14Run:
 15    uv run python examples/ports/balatro_feel/harness.py
 16"""
 17
 18from __future__ import annotations
 19
 20import sys
 21from collections.abc import Callable
 22from pathlib import Path
 23
 24_PORT_DIR = Path(__file__).resolve().parent
 25if str(_PORT_DIR) not in sys.path:
 26    sys.path.insert(0, str(_PORT_DIR))
 27
 28from main import HEIGHT, WIDTH, BalatroFeelRoot, CardTable  # noqa: E402
 29from nodes.card import SELECTION_LIFT  # noqa: E402
 30
 31from simvx.core import Key, MouseButton, Vec2  # noqa: E402
 32from simvx.core.testing.input_sim import InputSimulator  # noqa: E402
 33from simvx.graphics import App, save_png  # noqa: E402
 34
 35OUT = _PORT_DIR / "screenshots"
 36
 37
 38def capture_sequence() -> None:
 39    OUT.mkdir(exist_ok=True)
 40    app = App(width=WIDTH, height=HEIGHT, title="Balatro Feel Harness", visible=False)
 41    sim = InputSimulator()
 42    root = BalatroFeelRoot()
 43
 44    def hand():
 45        """The live HorizontalHandHolder, or None while the menu is up."""
 46        screen = root.screen
 47        return screen.hand if isinstance(screen, CardTable) else None
 48
 49    def card_centre(slot: int) -> Vec2:
 50        """Screen position of the card currently in `slot`, matching Card.contains."""
 51        holder = hand()
 52        card = next(c for c in holder.cards if c.slot_index == slot)
 53        return card.position + Vec2(0, -SELECTION_LIFT if card.selected else 0)
 54
 55    def hover(slot: int) -> Callable[[], None]:
 56        return lambda: sim.move_mouse(*card_centre(slot))
 57
 58    def press(slot: int) -> Callable[[], None]:
 59        return lambda: sim.press_mouse(MouseButton.LEFT, tuple(card_centre(slot)))
 60
 61    def release() -> Callable[[], None]:
 62        return lambda: sim.release_mouse(MouseButton.LEFT)
 63
 64    def drag(from_slot: int, to_slot: int, progress: float) -> Callable[[], None]:
 65        def _go() -> None:
 66            start, end = card_centre(from_slot), card_centre(to_slot)
 67            sim.move_mouse(start.x + (end.x - start.x) * progress, start.y - 18)
 68
 69        return _go
 70
 71    actions: dict[int, Callable[[], None]] = {}
 72    captures: dict[int, str] = {}
 73
 74    # Menu: the Play button is anchored dead centre, so its own rect is the target.
 75    captures[20] = "01_menu.png"
 76
 77    def press_play() -> None:
 78        x, y, w, h = root.screen.play_button.get_global_rect()
 79        sim.click((x + w * 0.5, y + h * 0.5))
 80
 81    actions[30] = press_play
 82    captures[70] = "02_idle.png"
 83
 84    # Hover the middle card, then click it to select.
 85    actions[80] = hover(3)
 86    captures[110] = "03_hover.png"
 87    actions[120] = press(3)
 88    actions[125] = release()
 89    captures[150] = "04_select.png"
 90
 91    # Press the leftmost card and walk it right across to slot 4.
 92    actions[170] = press(0)
 93    drag_start, drag_end = 175, 210
 94    for frame in range(drag_start, drag_end):
 95        actions[frame] = drag(0, 4, (frame - drag_start + 1) / (drag_end - drag_start))
 96    captures[205] = "05_drag.png"
 97    actions[220] = release()
 98    captures[245] = "06_dropped.png"
 99
100    # Select two more cards and play the hand with the keyboard shortcut.
101    actions[260] = press(2)
102    actions[265] = release()
103    actions[280] = press(5)
104    actions[285] = release()
105    actions[310] = lambda: sim.press_key(Key.SPACE)
106    actions[315] = lambda: sim.release_key(Key.SPACE)
107    captures[345] = "07_play.png"
108
109    def on_frame(idx: int, t: float) -> None:
110        action = actions.get(idx)
111        if action is not None:
112            action()
113
114    frames = app.run_headless(
115        root,
116        frames=380,
117        on_frame=on_frame,
118        capture_fn=lambda idx: idx in captures,
119    )
120
121    for frame_index, img in zip(sorted(captures), frames, strict=False):
122        out = OUT / captures[frame_index]
123        save_png(img, out)
124        print(f"saved {out}")
125    print(f"{len(frames)} screenshots written")
126
127
128if __name__ == "__main__":
129    capture_sequence()