Runner state setting and getting

tests/test_shortcuts.py

@@ -1,10 +1,12 @@
 import os
 from .context import YarnRunner
 
-compiled_yarn_f1 = open(os.path.join(os.path.dirname(
-    __file__), '../examples/yarn1/shortcuts.yarnc'), 'rb')
-names_csv_f1 = open(os.path.join(os.path.dirname(
-    __file__), '../examples/yarn1/shortcuts.csv'), 'r')
+compiled_yarn_fname1 = os.path.join(os.path.dirname(
+    __file__), '../examples/yarn1/shortcuts.yarnc')
+compiled_yarn_f1 = open(compiled_yarn_fname1, 'rb')
+names_csv_fname1 = os.path.join(os.path.dirname(
+    __file__), '../examples/yarn1/shortcuts.csv')
+names_csv_f1 = open(names_csv_fname1, 'r')
 compiled_yarn_f2 = open(os.path.join(os.path.dirname(
     __file__), '../examples/yarn2/shortcuts.yarnc'), 'rb')
 names_csv_f2 = open(os.path.join(os.path.dirname(
@@ -41,6 +43,62 @@ def test_shortcuts1():
     assert runner1.current_node == 'Start'
 
 
+def test_shortcuts2():
+    compiled_yarn_f3 = open(compiled_yarn_fname1, "rb")
+    names_csv_f3 = open(names_csv_fname1, "r")
+    runner3 = YarnRunner(compiled_yarn_f3, names_csv_f3)
+    assert "This is a test of shortcut functionality." == runner3.get_line()
+    assert not runner3.has_line()
+    assert not runner3.finished
+    runner3.choose(1)
+
+    assert "Option 2 selected." == runner3.get_line()
+    assert runner3.has_line()
+    assert "This is the last line." == runner3.get_line()
+    assert not runner3.has_line()
+    assert runner3.finished
+    assert runner3.current_node == 'Start'
+
+
+def test_shortcuts_json():
+    compiled_yarn_f3 = open(compiled_yarn_fname1, "rb")
+    names_csv_f3 = open(names_csv_fname1, "r")
+    runner3 = YarnRunner(compiled_yarn_f3, names_csv_f3)
+    assert "This is a test of shortcut functionality." == runner3.get_line()
+    assert not runner3.has_line()
+    assert not runner3.finished
+
+    choices = runner3.get_choices()
+
+    assert len(choices) == 4
+    assert choices[0]["text"] == "Option 1"
+    assert choices[1]["text"] == "Option 2"
+    assert choices[2]["text"] == "Option 3"
+    assert choices[3]["text"] == "Option 4"
+
+    dump = runner3.save()
+    compiled_yarn_f4 = open(compiled_yarn_fname1, "rb")
+    names_csv_f4 = open(names_csv_fname1, "r")
+    runner4 = YarnRunner(compiled_yarn_f4, names_csv_f4)
+    runner4.load(dump)
+
+    choices = runner4.get_choices()
+
+    assert len(choices) == 4
+    assert choices[0]["text"] == "Option 1"
+    assert choices[1]["text"] == "Option 2"
+    assert choices[2]["text"] == "Option 3"
+    assert choices[3]["text"] == "Option 4"
+
+    runner4.choose(0)
+
+    assert "Option 1 selected." == runner4.get_line()
+    assert runner4.has_line()
+    assert "This is the last line." == runner4.get_line()
+    assert not runner4.has_line()
+    assert runner4.finished
+    assert runner4.current_node == 'Start'
+
 def test_start_node_text2():
     assert "This is a test of shortcut functionality." == runner2.get_line()
     assert not runner2.has_line()

yarnrunner_python/runner.py

@@ -1,8 +1,11 @@
+from typing import Any, Dict, List, Optional, Union
+
 import csv
 import re
+import json
 from warnings import warn
 from google.protobuf import json_format
-from .yarn_spinner_pb2 import Program as YarnProgram, Instruction
+from .yarn_spinner_pb2 import Program as YarnProgram, Instruction  # type: ignore
 from .vm_std_lib import functions as std_lib_functions
 
 
@@ -33,6 +36,37 @@ def __init__(self, compiled_yarn_f, names_csv_f, autostart=True, enable_tracing=
         if autostart:
             self.resume()
 
+    def save(self) -> str:
+        dump = {
+            # "program_version": hash(self._compiled_yarn) + hash(self.string_lookup_table),
+            "visits": self.visits,
+            "variables": self.variables,
+            "current_node": self.current_node,
+            "line_buffer": self._line_buffer,
+            "option_buffer": self._option_buffer,
+            "vm_data_stack": self._vm_data_stack,
+            "vm_instruction_stack": [json.loads(json_format.MessageToJson(i)) for i in self._vm_instruction_stack],
+            "program_counter": self._program_counter,
+            "previous_instruction": json.loads(json_format.MessageToJson(self._previous_instruction)),
+            "finished": self.finished
+        }
+        return json.dumps(dump)
+
+    def load(self, data: str) -> None:
+        # if not data or "program_version" not in data or data["program_version"] != hash(self._compiled_yarn) + hash(self.string_lookup_table):
+        #     raise ValueError("Mismatched yarn version")
+        dump = json.loads(data)
+        self.visits = dump["visits"]
+        self.variables = dump["variables"]
+        self.current_node = dump["current_node"]
+        self._line_buffer = dump["line_buffer"]
+        self._option_buffer = dump["option_buffer"]
+        self._vm_data_stack = dump["vm_data_stack"]
+        self._vm_instruction_stack = [json_format.Parse(json.dumps(i), Instruction()) for i in dump["vm_instruction_stack"]]
+        self._program_counter = dump["program_counter"]
+        self._previous_instruction = json_format.Parse(json.dumps(dump["previous_instruction"]), Instruction())
+        self.finished = dump["finished"]
+
     def __construct_string_lookup_table(self):
         self.string_lookup_table = dict()
 
@@ -48,7 +82,7 @@ def resume(self):
         self.paused = False
         self.__process_instruction()
 
-    def __lookup_string(self, string_key):
+    def __lookup_string(self, string_key) -> str:
         if string_key not in self.string_lookup_table:
             raise Exception(
                 f"{string_key} is not a key in the string lookup table.")
@@ -242,7 +276,7 @@ def sanitize_quotes(arg):
             if type(ret) is str:
                 self._line_buffer.append(ret)
 
-    def __add_option(self, instruction):
+    def __add_option(self, instruction) -> None:
         title_string_key = instruction.operands[0].string_value
         choice_path = instruction.operands[1].string_value