Start to pylint/flake8 wizard cleanup

dwhswenson · Nov 5, 2021 · Nov 3, 2021 · Nov 3, 2021 · Nov 3, 2021 · Nov 4, 2021
commit 89bf2cbdb32c674a9bace925b544ceb6be13355b
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')
-