1"""Round-trip YAML editor for loop state action blocks.

2

3 Uses ruamel.yaml (round-trip mode) to preserve `action: |` block scalar

4 formatting when extracting or replacing a named state's action field.

5"""

6

7 from __future__ import annotations

8

9 from io import StringIO

10 from pathlib import Path

11

12 from ruamel.yaml import YAML

13 from ruamel.yaml.scalarstring import LiteralScalarString

14

15 from little_loops.file_utils import atomic_write

16

17

18 def extract_action(loop_yaml_path: Path, state_name: str) -> str:

19 """Return the action string for *state_name* in the loop YAML at *loop_yaml_path*.

20

21 Raises KeyError if *state_name* is not found in ``states``.

22 """

23 yaml = YAML(typ="rt")

24 data = yaml.load(loop_yaml_path)

25 return str(data["states"][state_name]["action"])

26

27

28 def replace_action(loop_yaml_path: Path, state_name: str, new_action: str) -> None:

29 """Replace the action for *state_name* in the loop YAML at *loop_yaml_path*.

30

31 Writes back in-place using atomic_write, preserving block scalar style

32 (``action: |``) for the modified state and leaving all other states and

33 keys unchanged.

34

35 Raises KeyError if *state_name* is not found in ``states``.

36 """

37 yaml = YAML(typ="rt")

38 data = yaml.load(loop_yaml_path)

39 # LiteralScalarString forces ruamel to emit `action: |` rather than

40 # choosing a quoted or flow scalar style for the new value.

41 data["states"][state_name]["action"] = LiteralScalarString(new_action)

42 buf = StringIO()

43 yaml.dump(data, buf)

44 atomic_write(loop_yaml_path, buf.getvalue())