move MockConsole, start to get_object decorator

dwhswenson · Jul 27, 2021 · Feb 18, 2021 · Feb 19, 2021 · Feb 27, 2021 · Mar 4, 2021
commit 4da1b9a2f34bd5dffb96ca2cec94019f8bcbc520
diff --git a/paths_cli/tests/wizard/mock_wizard.py b/paths_cli/tests/wizard/mock_wizard.py
@@ -10,3 +10,36 @@ def make_mock_retry_wizard(inputs):
     wizard = Wizard([])
     wizard.console.input = mock.Mock(side_effect=inputs)
     return wizard
+
+class MockConsole:
+    def __init__(self, inputs=None):
+        if isinstance(inputs, str):
+            inputs = [inputs]
+        elif inputs is None:
+            inputs = []
+        self.inputs = inputs
+        self._input_iter = iter(inputs)
+        self.log = []
+        self.width = 80
+        self.input_call_count = 0
+
+    def print(self, content=""):
+        self.log.append(content)
+
+    def input(self, content):
+        self.input_call_count += 1
+        user_input = next(self._input_iter)
+        self.log.append(content + " " + user_input)
+        return user_input
+
+    @property
+    def log_text(self):
+        return "\n".join(self.log)
+
+def mock_wizard(inputs):
+    wizard = Wizard([])
+    console = MockConsole(inputs)
+    wizard.console = console
+    return wizard
+
+
diff --git a/paths_cli/tests/wizard/test_wizard.py b/paths_cli/tests/wizard/test_wizard.py
@@ -1,7 +1,7 @@
 import pytest
 import mock
 from paths_cli.tests.wizard.mock_wizard import (
-    make_mock_wizard, make_mock_retry_wizard
+    MockConsole
 )
 import pathlib
 
@@ -10,30 +10,6 @@
 
 from paths_cli.wizard.wizard import *
 
-class MockConsole:
-    def __init__(self, inputs=None):
-        if isinstance(inputs, str):
-            inputs = [inputs]
-        elif inputs is None:
-            inputs = []
-        self.inputs = inputs
-        self._input_iter = iter(inputs)
-        self.log = []
-        self.width = 80
-        self.input_call_count = 0
-
-    def print(self, content=""):
-        self.log.append(content)
-
-    def input(self, content):
-        self.input_call_count += 1
-        user_input = next(self._input_iter)
-        self.log.append(content + " " + user_input)
-        return user_input
-
-    @property
-    def log_text(self):
-        return "\n".join(self.log)
 
 
 class TestWizard:

diff --git a/paths_cli/wizard/core.py b/paths_cli/wizard/core.py
@@ -64,3 +64,11 @@ def get_missing_object(wizard, obj_dict, display_name, fallback_func):
         obj = obj_dict[sel]
     return obj
 
+
+def get_object(func):
+    def inner(*args, **kwargs):
+        obj = None
+        while obj is None:
+            obj = func(*args, **kwargs)
+        return obj
+    return inner
diff --git a/paths_cli/wizard/openmm.py b/paths_cli/wizard/openmm.py
@@ -1,4 +1,5 @@
 from paths_cli.wizard.errors import FILE_LOADING_ERROR_MSG, not_installed
+from paths_cli.wizard.core import get_object
 try:
     from simtk import openmm as mm
     import mdtraj as md
@@ -16,6 +17,7 @@ def _openmm_serialization_helper(wizard, user_input):  # no-cov
                "simtk.openmm.openmm.XmlSerializer.html")
 
 
+@get_object
 def _load_openmm_xml(wizard, obj_type):
     xml = wizard.ask(
         f"Where is the XML file for your OpenMM {obj_type}?",
@@ -30,20 +32,17 @@ def _load_openmm_xml(wizard, obj_type):
     else:
         return obj
 
+@get_object
 def _load_topology(wizard):
     import openpathsampling as paths
-    topology = None
-    while topology is None:
-        filename = wizard.ask("Where is a PDB file describing your system?")
-        try:
-            snap = paths.engines.openmm.snapshot_from_pdb(filename)
-        except Exception as e:
-            wizard.exception(FILE_LOADING_ERROR_MSG, e)
-            continue
-
-        topology = snap.topology
+    filename = wizard.ask("Where is a PDB file describing your system?")
+    try:
+        snap = paths.engines.openmm.snapshot_from_pdb(filename)
+    except Exception as e:
+        wizard.exception(FILE_LOADING_ERROR_MSG, e)
+        return
 
-    return topology
+    return snap.topology
 
 def openmm(wizard):
     import openpathsampling as paths
@@ -56,17 +55,9 @@ def openmm(wizard):
                "To use OpenMM in OPS, you need to provide XML versions of "
                "your system, integrator, and some file containing "
                "topology information.")
-    system = integrator = topology = None
-    while system is None:
-        system = _load_openmm_xml(wizard, 'system')
-
-    while integrator is None:
-        integrator = _load_openmm_xml(wizard, 'integrator')
-
-    while topology is None:
-        topology = _load_topology(wizard)
-
-    n_frames_max = n_steps_per_frame = None
+    system = _load_openmm_xml(wizard, 'system')
+    integrator = _load_openmm_xml(wizard, 'integrator')
+    topology = _load_topology(wizard)
 
     n_steps_per_frame = wizard.ask_custom_eval(
         "How many MD steps per saved frame?", type_=int