first steps toward autodocumenting yaml structures

dwhswenson · Sep 27, 2021 · Mar 21, 2021 · Mar 21, 2021 · Mar 21, 2021 · Mar 21, 2021
commit 3597dd6ba69fc621bc2e835ebe825fe37df578ef
diff --git a/paths_cli/parsing/core.py b/paths_cli/parsing/core.py
@@ -1,4 +1,5 @@
 import os
+import json
 import importlib
 import yaml
 
@@ -22,14 +23,67 @@ def unlistify(obj, listified):
         obj = obj[0]
     return obj
 
+class Parameter:
+    SCHEMA = "http://openpathsampling.org/schemas/sim-setup/draft01.json"
+    def __init__(self, name, loader, required=True, json_type=None,
+                 description=None):
+        json_type = self._get_from_loader(loader, 'json_type', json_type)
+        description = self._get_from_loader(loader, 'description',
+                                            description)
+        if isinstance(json_type, str):
+            try:
+                json_type = json.loads(json_type)
+            except json.decoder.JSONDecodeError:
+                # TODO: maybe log this in case it represents an issue?
+                pass
+
+        self.name = name
+        self.loader = loader
+        self.json_type = json_type
+        self.description = description
+        self.required = required
+
+    @staticmethod
+    def _get_from_loader(loader, attr_name, attr):
+        if attr is None:
+            try:
+                attr = getattr(loader, attr_name)
+            except AttributeError:
+                pass
+        return attr
+
+
+    def __call__(self, *args, **kwargs):
+        # check correct call signature here
+        return self.loader(*args, **kwargs)
+
+    def to_json_schema(self, schema_context=None):
+        dct = {
+            'type': self.json_type,
+            'description': self.description,
+        }
+        return self.name, dct
+
+
 class InstanceBuilder:
-    # TODO: add schema as an input so we can autogenerate our JSON schema!
+    SCHEMA = "http://openpathsampling.org/schemas/sim-setup/draft01.json"
     def __init__(self, builder, attribute_table, optional_attributes=None,
-                 defaults=None, module=None, remapper=None):
+                 defaults=None, module=None, remapper=None, parameters=None,
+                 object_type=None, name=None):
+        # temporary apporach to override attribute_table
+        if attribute_table is None and parameters is not None:
+            attribute_table = {p.name: p.loader for p in parameters
+                               if p.required}
+        if optional_attributes is None and parameters is not None:
+            optional_attributes = {p.name: p.loader for p in parameters
+                                   if not p.required}
+        self.object_type = object_type
+        self.name = name
         self.module = module
         self.builder = builder
         self.builder_name = str(self.builder)
         self.attribute_table = attribute_table
+        self.parameters = parameters
         if optional_attributes is None:
             optional_attributes = {}
         self.optional_attributes = optional_attributes
@@ -42,6 +96,29 @@ def __init__(self, builder, attribute_table, optional_attributes=None,
         self.defaults = defaults
         self.logger = logging.getLogger(f"parser.InstanceBuilder.{builder}")
 
+    @property
+    def schema_name(self):
+        if not self.name.endswith(self.object_type):
+            schema_name = f"{self.name}-{self.object_type}"
+        else:
+            schema_name = name
+        return schema_name
+
+    def to_json_schema(self, schema_context=None):
+        parameters = dict(p.to_json_schema() for p in self.parameters)
+        properties = {
+            'name': {'type': 'string'},
+            'type': {'type': 'string',
+                     'enum': [self.name]},
+        }
+        properties.update(parameters)
+        required = [p.name for p in self.parameters if p.required]
+        dct = {
+            'properties': properties,
+            'required': required,
+        }
+        return self.schema_name, dct
+
     def select_builder(self, dct):
         if self.module is not None:
             builder = getattr(importlib.import_module(self.module), self.builder)

diff --git a/paths_cli/parsing/cvs.py b/paths_cli/parsing/cvs.py
@@ -1,7 +1,7 @@
 import os
 import importlib
 
-from .core import Parser, InstanceBuilder, custom_eval
+from .core import Parser, InstanceBuilder, custom_eval, Parameter
 from .topology import build_topology
 from .errors import InputError
 
@@ -27,33 +27,31 @@ def cv_prepare_dict(dct):
 # MDTraj-specific
 
 mdtraj_source = AllowedPackageHandler("mdtraj")
-MDTRAJ_ATTRS = {
-    'topology': build_topology,
-    'func': mdtraj_source,
-    'kwargs': lambda kwargs: {key: custom_eval(arg)
-                              for key, arg in kwargs.items()},
-}
+# MDTRAJ_ATTRS = {
+    # 'topology': build_topology,
+    # 'func': mdtraj_source,
+    # 'kwargs': lambda kwargs: {key: custom_eval(arg)
+                              # for key, arg in kwargs.items()},
+# }
 
-MDTRAJ_SCHEMA = {
-    'topology': {
-        'type': 'string',
-        'description': 'topology from file or engine name',
-    },
-    'func': {
-        'type': 'string',
-        'description': 'MDTraj function, e.g., ``compute_distances``',
-    },
-    'kwargs': {
-        'type': 'object',
-        'description': 'keyword arguments for ``func``',
-    },
-}
+MDTRAJ_PARAMETERS = [
+    Parameter('topology', build_topology),
+    Parameter('func', mdtraj_source, json_type='string',
+              description="MDTraj function, e.g., ``compute_distances``"),
+    Parameter('kwargs', lambda kwargs: {key: custom_eval(arg)
+                                        for key, arg in kwargs.items()},
+              json_type='object',
+              description="keyword arguments for ``func``",
+              required=True),  # TODO: bug in current: shoudl be req=False
+]
 
 build_mdtraj_function_cv = InstanceBuilder(
+    attribute_table=None,  # temp
     module='openpathsampling.experimental.storage.collective_variables',
     builder='MDTrajFunctionCV',
-    attribute_table=MDTRAJ_ATTRS,
-    remapper = cv_prepare_dict,
+    # attribute_table=MDTRAJ_ATTRS,
+    parameters=MDTRAJ_PARAMETERS,
+    remapper=cv_prepare_dict,
 )
 
 # TODO: this should replace TYPE_MAPPING and cv_parser

diff --git a/paths_cli/parsing/topology.py b/paths_cli/parsing/topology.py
@@ -24,9 +24,11 @@ def get_topology_from_file(dct):
 
 class MultiStrategyBuilder:
     # move to core
-    def __init__(self, strategies, label):
+    def __init__(self, strategies, label, description=None, json_type=None):
         self.strategies = strategies
         self.label = label
+        self.description = description
+        self.json_type = json_type
 
     def __call__(self, dct):
         for strategy in self.strategies:
@@ -37,6 +39,9 @@ def __call__(self, dct):
         # only get here if we failed
         raise InputError.invalid_input(dct, self.label)
 
-build_topology = MultiStrategyBuilder([get_topology_from_file,
-                                       get_topology_from_engine],
-                                      label='topology')
+build_topology = MultiStrategyBuilder(
+    [get_topology_from_file, get_topology_from_engine],
+    label='topology',
+    description="topology from file or engine name",
+    json_type='string'
+)
diff --git a/paths_cli/tests/conftest.py b/paths_cli/tests/conftest.py
@@ -21,14 +21,17 @@ def cv_and_states():
     state_B = paths.CVDefinedVolume(cv, 1, float("inf")).named("B")
     return cv, state_A, state_B
 
+@pytest.fixture
+def transition_traj():
+    init_traj = make_1d_traj([-0.1, 0.1, 0.3, 0.5, 0.7, 0.9, 1.1],
+                             [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0])
+    return init_traj
 
 @pytest.fixture
-def tps_network_and_traj(cv_and_states):
+def tps_network_and_traj(cv_and_states, transition_traj):
     _, state_A, state_B = cv_and_states
     network = paths.TPSNetwork(state_A, state_B)
-    init_traj = make_1d_traj([-0.1, 0.1, 0.3, 0.5, 0.7, 0.9, 1.1],
-                             [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0])
-    return (network, init_traj)
+    return network, transition_traj
 
 @pytest.fixture
 def tps_fixture(flat_engine, tps_network_and_traj):
@@ -38,3 +41,12 @@ def tps_fixture(flat_engine, tps_network_and_traj):
                                             engine=flat_engine)
     init_conds = scheme.initial_conditions_from_trajectories(traj)
     return (scheme, network, flat_engine, init_conds)
+
+
+@pytest.fixture
+def tis_network(cv_and_states):
+    cv, state_A, state_B = cv_and_states
+    interfaces = paths.VolumeInterfaceSet(cv, float("-inf"),
+                                          [0.0, 0.1, 0.2])
+    network = paths.MISTISNetwork([(state_A, interfaces, state_B)])
+    return network