openpathsampling · dwhswenson · Nov 5, 2021 · Nov 3, 2021 · Nov 3, 2021 · Nov 3, 2021
diff --git a/paths_cli/tests/wizard/test_load_from_ops.py b/paths_cli/tests/wizard/test_load_from_ops.py
@@ -6,7 +6,7 @@
 from paths_cli.tests.wizard.mock_wizard import mock_wizard
 
 from paths_cli.wizard.load_from_ops import (
-    named_objs_helper, _get_ops_storage, _get_ops_object, load_from_ops
+    _get_ops_storage, _get_ops_object, load_from_ops
 )
 
 # for some reason I couldn't get these to work with MagicMock
@@ -44,12 +44,6 @@ def ops_file_fixture():
     storage = FakeStorage(foo)
     return storage
 
-def test_named_objs_helper(ops_file_fixture):
-    helper_func = named_objs_helper(ops_file_fixture, 'foo')
-    result = helper_func('any')
-    assert "what I found" in result
-    assert "bar" in result
-    assert "baz" in result
 
 @pytest.mark.parametrize('with_failure', [False, True])
 def test_get_ops_storage(tmpdir, with_failure):

diff --git a/paths_cli/wizard/core.py b/paths_cli/wizard/core.py
@@ -1,12 +1,17 @@
-import random
-from paths_cli.wizard.tools import a_an
-
-from collections import namedtuple
+def interpret_req(req):
+    """Create user-facing string representation of the input requirement.
 
-WIZARD_STORE_NAMES = ['engines', 'cvs', 'states', 'networks', 'schemes']
-WizardSay = namedtuple("WizardSay", ['msg', 'mode'])
+    Parameters
+    ----------
+    req : Tuple[..., int, int]
+        req[1] is the minimum number of objects to create; req[2] is the
+        maximum number of objects to create
 
-def interpret_req(req):
+    Returns
+    -------
+    str :
+        human-reading string for how many objects to create
+    """
     _, min_, max_ = req
     string = ""
     if min_ == max_:
@@ -23,7 +28,32 @@ def interpret_req(req):
     return string
 
 
+# TODO: REFACTOR: It looks like get_missing_object may be redundant with
+# other code for obtaining prerequisites for a function
 def get_missing_object(wizard, obj_dict, display_name, fallback_func):
+    """Get a prerequisite object.
+
+    The ``obj_dict`` here is typically a mapping of objects known by the
+    Wizard. If it is empty, the ``fallback_func`` is used to create a new
+    object. If it has exactly 1 entry, that is used implicitly. If it has
+    more than 1 entry, the user must select which one to use.
+
+    Parameters
+    ----------
+    wizard : :class:`.Wizard`
+        the wizard for user interaction
+    obj_dict : Dict[str, Any]
+        mapping of object name to object
+    display_name: str
+        the user-facing name of this type of object
+    fallback_func: Callable[:class:`.Wizard`] -> Any
+        method to create a new object of this type
+
+    Returns
+    -------
+    Any :
+        the prerequisite object
+    """
     if len(obj_dict) == 0:
         obj = fallback_func(wizard)
     elif len(obj_dict) == 1:
@@ -37,10 +67,22 @@ def get_missing_object(wizard, obj_dict, display_name, fallback_func):
 
 
 def get_object(func):
+    """Decorator to wrap methods for obtaining objects from user input.
+
+    This decorator implements the user interaction loop when dealing with a
+    single user input. The wrapped function is intended to create some
+    object. If the user's input cannot create a valid object, the wrapped
+    function should return None.
+
+    Parameters
+    ----------
+    func : Callable
+        object creation method to wrap; should return None on failure
+    """
+    # TODO: use functools.wraps?
     def inner(*args, **kwargs):
         obj = None
         while obj is None:
             obj = func(*args, **kwargs)
         return obj
     return inner
-
diff --git a/paths_cli/wizard/cvs.py b/paths_cli/wizard/cvs.py
@@ -1,14 +1,14 @@
+from functools import partial
+from collections import namedtuple
+import numpy as np
+
 from paths_cli.compiling.tools import mdtraj_parse_atomlist
 from paths_cli.wizard.plugin_classes import (
     LoadFromOPS, WizardObjectPlugin, WrapCategory
 )
 from paths_cli.wizard.core import get_object
 import paths_cli.wizard.engines
 
-from functools import partial
-from collections import namedtuple
-import numpy as np
-
 from paths_cli.wizard.parameters import (
     FromWizardPrerequisite
 )
@@ -27,6 +27,10 @@
     "You should specify atom indices enclosed in double brackets, e.g. "
     "[{list_range_natoms}]"
 )
+_MDTrajParams = namedtuple("_MDTrajParams", ['period', 'n_atoms',
+                                             'kwarg_name', 'cv_user_str'])
+_MDTRAJ_INTRO = "We'll make a CV that measures the {user_str}."
+
 
 # TODO: implement so the following can be the help string:
 # _ATOM_INDICES_HELP_STR = (
@@ -56,6 +60,25 @@
 
 @get_object
 def _get_atom_indices(wizard, topology, n_atoms, cv_user_str):
+    """Parameter loader for atom_indices parameters in MDTraj.
+
+    Parameters
+    ----------
+    wizard : :class:`.Wizard`
+        wizard for user interaction
+    topology :
+        topology (reserved for future use)
+    n_atoms : int
+        number of atoms to define this CV (i.e., 2 for a distance; 3 for an
+        angle; 4 for a dihedral)
+    cv_user_str : str
+        user-facing name for the CV being created
+
+    Returns
+    -------
+    :class`np.ndarray` :
+        array of indices for the MDTraj function
+    """
     helper = Helper(_ATOM_INDICES_HELP_STR.format(
         list_range_natoms=list(range(n_atoms))
     ))
@@ -67,14 +90,14 @@ def _get_atom_indices(wizard, topology, n_atoms, cv_user_str):
     except Exception as e:
         wizard.exception(f"Sorry, I didn't understand '{atoms_str}'.", e)
         helper("?")
-        return
+        return None
 
     return arr
 
-_MDTrajParams = namedtuple("_MDTrajParams", ['period', 'n_atoms',
-                                             'kwarg_name', 'cv_user_str'])
 
 def _mdtraj_cv_builder(wizard, prereqs, func_name):
+    """General function to handle building MDTraj CVs.
+    """
     from openpathsampling.experimental.storage.collective_variables import \
             MDTrajFunctionCV
     dct = TOPOLOGY_CV_PREREQ(wizard)
@@ -108,9 +131,9 @@ def _mdtraj_cv_builder(wizard, prereqs, func_name):
     return MDTrajFunctionCV(func, topology, period_min=period_min,
                             period_max=period_max, **kwargs)
 
-_MDTRAJ_INTRO = "We'll make a CV that measures the {user_str}."
 
 def _mdtraj_summary(wizard, context, result):
+    """Standard summary of MDTraj CVs: function, atom, topology"""
     cv = result
     func = cv.func
     topology = cv.topology
@@ -121,6 +144,7 @@ def _mdtraj_summary(wizard, context, result):
                f"  Topology: {repr(topology.mdtraj)}")
     return [summary]
 
+
 if HAS_MDTRAJ:
     MDTRAJ_DISTANCE = WizardObjectPlugin(
         name='Distance',
@@ -152,10 +176,24 @@ def _mdtraj_summary(wizard, context, result):
                      "four atoms"),
         summary=_mdtraj_summary,
     )
-
     # TODO: add RMSD -- need to figure out how to select a frame
 
+
 def coordinate(wizard, prereqs=None):
+    """Builder for coordinate CV.
+
+    Parameters
+    ----------
+    wizard : :class:`.Wizard`
+        wizard for user interaction
+    prereqs :
+        prerequisites (unused in this method)
+
+    Return
+    ------
+    CoordinateFunctionCV :
+        the OpenPathSampling CV for this selecting this coordinate
+    """
     # TODO: atom_index should be from wizard.ask_custom_eval
     from openpathsampling.experimental.storage.collective_variables import \
             CoordinateFunctionCV
@@ -174,12 +212,13 @@ def coordinate(wizard, prereqs=None):
                          f"atom {atom_index}?")
         try:
             coord = {'x': 0, 'y': 1, 'z': 2}[xyz]
-        except KeyError as e:
+        except KeyError:
             wizard.bad_input("Please select one of 'x', 'y', or 'z'")
 
     cv = CoordinateFunctionCV(lambda snap: snap.xyz[atom_index][coord])
     return cv
 
+
 COORDINATE_CV = WizardObjectPlugin(
     name="Coordinate",
     category="cv",
@@ -202,6 +241,7 @@ def coordinate(wizard, prereqs=None):
             "you can also create your own and load it from a file.")
 )
 
+
 if __name__ == "__main__":  # no-cov
     from paths_cli.wizard.run_module import run_category
     run_category('cv')
diff --git a/paths_cli/wizard/engines.py b/paths_cli/wizard/engines.py
@@ -1,6 +1,4 @@
-import paths_cli.wizard.openmm as openmm
 from paths_cli.wizard.plugin_classes import LoadFromOPS, WrapCategory
-from functools import partial
 
 _ENGINE_HELP = "An engine describes how you'll do the actual dynamics."
 ENGINE_PLUGIN = WrapCategory(
@@ -17,4 +15,3 @@
 if __name__ == "__main__":  # no-cov
     from paths_cli.wizard.run_module import run_category
     run_category('engine')
-
diff --git a/paths_cli/wizard/errors.py b/paths_cli/wizard/errors.py
@@ -1,23 +1,42 @@
 class ImpossibleError(Exception):
+    """Error to throw for sections that should be unreachable code"""
     def __init__(self, msg=None):
         if msg is None:
             msg = "Something went really wrong. You should never see this."
         super().__init__(msg)
 
+
 class RestartObjectException(BaseException):
+    """Exception to indicate the restart of an object.
+
+    Raised when the user issues a command to cause an object restart.
+    """
     pass
 
+
 def not_installed(wizard, package, obj_type):
+    """Behavior when an integration is not installed.
+
+    In actual practice, this calling code should ensure this doesn't get
+    used. However, we keep it around as a defensive practice.
+
+    Parameters
+    ----------
+    package : str
+        name of the missing package
+    obj_type : str
+        name of the object category that would have been created
+    """
     retry = wizard.ask(f"Hey, it looks like you don't have {package} "
                        "installed. Do you want to try a different "
                        f"{obj_type}, or do you want to quit?",
-                       options=["[R]etry", "[Q]uit"])
+                       options=["[r]etry", "[q]uit"])
     if retry == 'r':
         raise RestartObjectException()
-    elif retry == 'q':
+    if retry == 'q':
+        # TODO: maybe raise QuitWizard instead?
         exit()
-    else:  # no-cov
-        raise ImpossibleError()
+    raise ImpossibleError()  # -no-cov-
 
 
 FILE_LOADING_ERROR_MSG = ("Sorry, something went wrong when loading that "

diff --git a/paths_cli/wizard/helper.py b/paths_cli/wizard/helper.py
@@ -1,24 +1,31 @@
+import sys
 from .errors import RestartObjectException
 
 
 class QuitWizard(BaseException):
-    pass
+    """Exception raised when user expresses desire to quit the wizard"""
 
 
 # the following command functions take cmd and ctx -- future commands might
 # use the full command text or the context internally.
 
 def raise_quit(cmd, ctx):
+    """Command function to quit the wizard (with option to save).
+    """
     raise QuitWizard()
 
 
 def raise_restart(cmd, ctx):
+    """Command function to restart the current object.
+    """
     raise RestartObjectException()
 
 
 def force_exit(cmd, ctx):
+    """Command function to force immediate exit.
+    """
     print("Exiting...")
-    exit()
+    sys.exit()
 
 
 HELPER_COMMANDS = {

diff --git a/paths_cli/wizard/joke.py b/paths_cli/wizard/joke.py
@@ -16,31 +16,38 @@
     "It would also be a good name for a death metal band.",
 ]
 
+
 def _joke1(name, obj_type):  # no-cov
     return (f"I probably would have named it something like "
             f"'{random.choice(_NAMES)}'.")
 
+
 def _joke2(name, obj_type):  # no-cov
     thing = random.choice(_THINGS)
     joke = (f"I had {a_an(thing)} {thing} named '{name}' "
             f"when I was young.")
     return joke
 
+
 def _joke3(name, obj_type):  # no-cov
     return (f"I wanted to name my {random.choice(_SPAWN)} '{name}', but my "
             f"wife wouldn't let me.")
 
+
 def _joke4(name, obj_type):  # no-cov
     a_an_thing = a_an(obj_type) + f" {obj_type}"
     return random.choice(_MISC).format(name=name, obj_type=obj_type,
                                        a_an_thing=a_an_thing)
 
+
 def name_joke(name, obj_type):  # no-cov
+    """Make a joke about the naming process."""
     jokes = [_joke1, _joke2, _joke3, _joke4]
     weights = [5, 5, 3, 7]
     joke = random.choices(jokes, weights=weights)[0]
     return joke(name, obj_type)
 
+
 if __name__ == "__main__":  # no-cov
     for _ in range(5):
         print()