harness.py

Part of Mr. Rescue.

  1"""Scripted headless playthrough: the implementation behind ``main.py --test``.
  2
  3:func:`capture_playthrough` drives a real :class:`MrRescueRoot` with
  4:class:`~simvx.core.InputSimulator` and writes one PNG per stage into
  5``screenshots/``.
  6
  7The pilot below plays from live world state rather than from a fixed key
  8timeline: it presses "up" on the frame it is actually standing on a ladder,
  9steps off once it has climbed a full storey, and presses "grab" on the frame a
 10civilian comes into reach. That is what makes each screenshot show the outcome
 11its name promises instead of the key that was sent.
 12
 13Module-level ``random`` is seeded first: civilians, fire bugs and screenshake all
 14draw from it, so without a seed the run would not repeat frame for frame and the
 15capture points would drift.
 16
 17Run it through the port's own entry point::
 18
 19    uv run python examples/ports/mr_rescue/main.py --test
 20"""
 21
 22from __future__ import annotations
 23
 24import random
 25from collections.abc import Callable
 26from pathlib import Path
 27
 28from nodes.building import FLOOR_HEIGHT
 29from nodes.tile_grid import TILE_SIZE
 30
 31from simvx.core import InputSimulator, Key, Node
 32from simvx.graphics import App, save_png
 33
 34# Height of one storey in pixels: the exact distance a ladder spans.
 35STOREY_PX = FLOOR_HEIGHT * TILE_SIZE
 36
 37RUN_SEED = 42  # scene generator + module-level ``random``
 38START_FRAME = 30  # confirm on the title screen
 39RUN_FRAME = 50  # start running right
 40SPRAY_FRAME = 70  # open the hose
 41SPRAY_STOP_X = 400.0  # close it again short of the ladder, so climbing can start
 42CARRY_WALK = 40  # frames spent carrying the civilian clear of the group
 43CARRY_HOLD = 60  # frames spent carrying before the suit is pushed over
 44RESTART_DELAY = 40  # frames on the end screen before confirming
 45
 46# Frame at which each screenshot is taken. The pilot is state-driven but the run
 47# is deterministic, so these land on the stages listed in ``capture_playthrough``.
 48CAPTURES = {
 49    20: "01_title.png",
 50    48: "02_world_overview.png",
 51    157: "03_water_spray.png",
 52    330: "04_ladder_climb.png",
 53    500: "05_upper_floor.png",
 54    580: "06_civilian_grab.png",
 55    620: "07_end_screen.png",
 56    664: "08_post_restart.png",
 57}
 58TOTAL_FRAMES = 700
 59
 60
 61class _Pilot:
 62    """Plays one run, deciding each frame's input from the live scene."""
 63
 64    def __init__(self, sim: InputSimulator):
 65        self._sim = sim
 66        self._held: set[Key] = set()
 67        self._taps: dict[Key, int] = {}
 68        self.stage = "title"
 69        self._mount_y = 0.0
 70        self._carry_frame = 0
 71        self._end_frame = -1
 72
 73    # ------------------------------------------------------------- keys
 74
 75    def _hold(self, key: Key):
 76        if key not in self._held:
 77            self._sim.press_key(key)
 78            self._held.add(key)
 79
 80    def _let_go(self, key: Key):
 81        if key in self._held:
 82            self._sim.release_key(key)
 83            self._held.discard(key)
 84
 85    def _tap(self, idx: int, key: Key):
 86        """Press *key* now and release it two frames later (a one-shot action)."""
 87        self._hold(key)
 88        self._taps[key] = idx + 2
 89
 90    # ------------------------------------------------------------- driver
 91
 92    def step(self, idx: int, root: Node):
 93        for key, due in list(self._taps.items()):
 94            if idx >= due:
 95                self._let_go(key)
 96                del self._taps[key]
 97
 98        if root.phase == "end":
 99            self._end(idx)
100            return
101        if root.phase == "menu":
102            if self.stage == "title" and idx >= START_FRAME:
103                self._tap(idx, Key.ENTER)
104                self.stage = "advance"
105            return
106
107        scene = root.scene
108        if self.stage == "advance":
109            self._advance(idx, scene)
110        elif self.stage == "climb":
111            self._climb(scene.player)
112        elif self.stage == "cross":
113            self._cross(idx, scene.player)
114        elif self.stage == "carry":
115            self._carry(idx, scene.player)
116
117    # ------------------------------------------------------------- stages
118
119    def _advance(self, idx: int, scene):
120        """Run right along the ground floor with the hose open, up to the ladder."""
121        player, grid = scene.player, scene.grid
122        if idx >= RUN_FRAME:
123            self._hold(Key.RIGHT)
124        if SPRAY_FRAME <= idx and player.position.x < SPRAY_STOP_X:
125            self._hold(Key.LEFT_SHIFT)
126        else:
127            self._let_go(Key.LEFT_SHIFT)
128        # Mount the moment the gun-height cell is a rung: holding "up" over a
129        # ladder is exactly how a player climbs.
130        if grid.is_ladder(*grid.cell_for(player.position.x, player.position.y - 11)):
131            self._let_go(Key.RIGHT)
132            self._hold(Key.UP)
133            self._mount_y = float(player.position.y)
134            self.stage = "climb"
135
136    def _climb(self, player):
137        """Hold "up" for one full storey, then step off sideways onto the floor."""
138        if player.position.y <= self._mount_y - STOREY_PX:
139            self._let_go(Key.UP)
140            self._hold(Key.LEFT)
141            self.stage = "cross"
142
143    def _cross(self, idx: int, player):
144        """Walk the upper floor hosing the fires down, and grab the first civilian."""
145        self._hold(Key.LEFT)
146        self._hold(Key.LEFT_SHIFT)
147        if player.can_grab:
148            self._let_go(Key.LEFT)
149            self._let_go(Key.LEFT_SHIFT)
150            self._tap(idx, Key.E)
151            self._carry_frame = idx
152            self.stage = "carry"
153
154    def _carry(self, idx: int, player):
155        """Carry her back the way we came, then let the suit cook."""
156        held = idx - self._carry_frame
157        if held < CARRY_WALK:
158            self._hold(Key.RIGHT)
159        else:
160            self._let_go(Key.RIGHT)
161        # Push the suit past its heat limit so the end screen is reachable inside
162        # the frame budget; real play gets there through fire exposure.
163        if held >= CARRY_HOLD:
164            player.temperature = player.max_temperature + 1.0
165
166    def _end(self, idx: int):
167        if self._end_frame < 0:
168            self._end_frame = idx
169            for key in list(self._held):
170                self._let_go(key)
171        elif idx == self._end_frame + RESTART_DELAY:
172            self._tap(idx, Key.ENTER)
173
174
175def capture_playthrough(
176    root_factory: Callable[..., Node],
177    *,
178    window_size: tuple[int, int],
179    bg_colour,
180    out_dir: Path,
181) -> None:
182    """Run the scripted playthrough headlessly and save the stage screenshots.
183
184    The stages, in the order they are captured:
185
186    ``01_title``
187        the title screen.
188    ``02_world_overview``
189        the generated building, camera on the player at his ground-floor spawn.
190    ``03_water_spray``
191        running right with the hose open, the jet landing on a burning
192        ground-floor fire (a second one is already out: score 20, fires 5).
193    ``04_ladder_climb``
194        on the ladder, part-way between the ground floor and the one above.
195    ``05_upper_floor``
196        stepped off the ladder and crossing the floor above, hosing down the two
197        fires that were burning on it.
198    ``06_civilian_grab``
199        carrying a civilian, picked up on the frame she came into reach.
200    ``07_end_screen``
201        the overheat end screen with the run summary.
202    ``08_post_restart``
203        back on the title screen after confirming a restart.
204    """
205    out_dir.mkdir(exist_ok=True)
206    random.seed(RUN_SEED)
207
208    width, height = window_size
209    app = App(title="Mr. Rescue (test)", width=width, height=height, visible=False, bg_colour=bg_colour)
210    sim = InputSimulator()
211    root = root_factory(seed=RUN_SEED)
212    pilot = _Pilot(sim)
213
214    def drive(idx: int, _t: float):
215        pilot.step(idx, root)
216        return None
217
218    capture_frames = sorted(CAPTURES)
219    frames = app.run_headless(
220        root,
221        frames=TOTAL_FRAMES,
222        on_frame=drive,
223        capture_frames=capture_frames,
224    )
225    for idx, frame in zip(capture_frames, frames, strict=False):
226        if frame is None:
227            continue
228        # Headless captures come back with the scene's own alpha; force opaque
229        # so the PNGs composite predictably in docs and image viewers.
230        frame[..., 3] = 255
231        path = out_dir / CAPTURES[idx]
232        save_png(frame, path)
233        print(f"saved {path}")