Merge pull request #63 from dwhswenson/wizard-quit-save

Wizard: `!quit` command should give option to save

Author: dwhswenson

paths_cli/tests/wizard/test_wizard.py

@@ -368,3 +368,31 @@ def test_run_wizard(self, toy_engine):
         assert len(storage.networks) == len(storage.schemes) == 0
         assert len(storage.engines) == 1
         assert storage.engines[toy_engine.name] == toy_engine
+
+    def test_run_wizard_quit(self):
+        console = MockConsole()
+        self.wizard.console = console
+        self.wizard._patched = True
+        step = mock.Mock(
+            func=mock.Mock(side_effect=QuitWizard()),
+            display_name='Engine',
+            store_name='engines',
+            minimum=1,
+            maximum=1
+        )
+        self.wizard.steps = [step]
+        mock_ask_save = mock.Mock(return_value=False)
+        with mock.patch.object(Wizard, '_ask_save', mock_ask_save):
+            self.wizard.run_wizard()
+        assert "Goodbye!" in self.wizard.console.log_text
+
+    @pytest.mark.parametrize('inputs', ['y', 'n'])
+    def test_ask_save(self, inputs):
+        expected = {'y': True, 'n': False}[inputs]
+        console = MockConsole(['foo', inputs])
+        self.wizard.console = console
+        result = self.wizard._ask_save()
+        assert result is expected
+        assert "Before quitting" in self.wizard.console.log_text
+        assert "Sorry" in self.wizard.console.log_text
+

paths_cli/wizard/wizard.py

@@ -9,7 +9,7 @@
     FILE_LOADING_ERROR_MSG, RestartObjectException
 )
 from paths_cli.wizard.joke import name_joke
-from paths_cli.wizard.helper import Helper
+from paths_cli.wizard.helper import Helper, QuitWizard
 from paths_cli.compiling.tools import custom_eval
 
 import shutil
@@ -313,13 +313,33 @@ def run_wizard(self):
         # TODO: next line is only temporary
         self.say("Today I'll help you set up a 2-state TPS simulation.")
         self._patch()  # try to hide the slowness of our first import
-        for step in self.steps:
-            req = step.store_name, step.minimum, step.maximum
-            do_another = True
-            while do_another:
-                do_another = self._do_one(step, req)
-        storage = self.get_storage()
-        self.save_to_file(storage)
+        try:
+            for step in self.steps:
+                req = step.store_name, step.minimum, step.maximum
+                do_another = True
+                while do_another:
+                    do_another = self._do_one(step, req)
+        except QuitWizard:
+            do_save = self._ask_save()
+        else:
+            do_save = True
+
+        if do_save:
+            storage = self.get_storage()
+            self.save_to_file(storage)
+        else:
+            self.say("Goodbye! 👋")
+
+    @get_object
+    def _ask_save(self):
+        do_save_char = self.ask("Before quitting, would you like to save "
+                                "the objects you've created so far?")
+        try:
+            do_save = yes_no(do_save_char)
+        except Exception:
+            self.bad_input("Sorry, I didn't understance that.")
+            return None
+        return do_save
 
 
 # FIXED_LENGTH_TPS_WIZARD