openpathsampling · dwhswenson · Jan 1, 2021 · Oct 23, 2020 · Oct 23, 2020 · Nov 27, 2020
diff --git a/devtools/tests_require.txt b/devtools/tests_require.txt
@@ -3,3 +3,6 @@ nose
 pytest
 pytest-cov
 coveralls
+# following are for SimStore integration
+dill
+sqlalchemy
diff --git a/paths_cli/param_core.py b/paths_cli/param_core.py
@@ -72,22 +72,48 @@ class StorageLoader(AbstractLoader):
     mode : 'r', 'w', or 'a'
         the mode for the file
     """
+    has_simstore_patch = False
     def __init__(self, param, mode):
         super(StorageLoader, self).__init__(param)
         self.mode = mode
 
+    @staticmethod
+    def _is_simstore(name):
+        return name.endswith(".db") or name.endswith(".sql")
+
     def _workaround(self, name):
         # this is messed up... for some reason, storage doesn't create a new
         # file in append mode. That may be a bug
         import openpathsampling as paths
-        if self.mode == 'a' and not os.path.exists(name):
+        needs_workaround = (
+            self.mode == 'a'
+            and not os.path.exists(name)
+            and not self._is_simstore(name)
+        )
+        if needs_workaround:
             st = paths.Storage(name, mode='w')
             st.close()
 
     def get(self, name):
-        import openpathsampling as paths
-        self._workaround(name)
-        return paths.Storage(name, mode=self.mode)
+        if self._is_simstore(name):
+            import openpathsampling as paths
+            from openpathsampling.experimental.storage import \
+                    Storage, monkey_patch_all
+
+            if not self.has_simstore_patch:
+                paths = monkey_patch_all(paths)
+                paths.InterfaceSet.simstore = True
+                StorageLoader.has_simstore_patch = True
+
+            from openpathsampling.experimental.simstore import \
+                SQLStorageBackend
+            backend = SQLStorageBackend(name, mode=self.mode)
+            storage = Storage.from_backend(backend)
+        else:
+            from openpathsampling import Storage
+            self._workaround(name)
+            storage = Storage(name, self.mode)
+        return storage
 
 
 class OPSStorageLoadNames(AbstractLoader):
@@ -200,10 +226,20 @@ class GetOnlySnapshot(Getter):
     def __init__(self, store_name="snapshots"):
         super().__init__(store_name)
 
+    def _min_num_snapshots(self, storage):
+        # For netcdfplus, we see 2 snapshots when there is only one
+        # (reversed copy gets saved). For SimStore, we see only one.
+        import openpathsampling as paths
+        if isinstance(storage, paths.netcdfplus.NetCDFPlus):
+            min_snaps = 2
+        else:
+            min_snaps = 1
+        return min_snaps
+
     def __call__(self, storage):
         store = getattr(storage, self.store_name)
-        if len(store) == 2:
-            # this is really only 1 snapshot; reversed copy gets saved
+        min_snaps = self._min_num_snapshots(storage)
+        if len(store) == min_snaps:
             return store[0]
 
 
@@ -270,7 +306,14 @@ def get(self, storage, name):
             result = _try_strategies(self.none_strategies, storage)
 
         if result is None:
-            raise RuntimeError("Couldn't find %s", name)
+            if name is None:
+                msg = "Couldn't guess which item to use from " + self.store
+            else:
+                msg = "Couldn't find {name} is {store}".format(
+                    name=name,
+                    store=self.store
+                )
+            raise RuntimeError(msg)
 
         return result
 

diff --git a/paths_cli/tests/commands/test_append.py b/paths_cli/tests/commands/test_append.py
@@ -8,22 +8,22 @@
 import openpathsampling as paths
 
 def make_input_file(tps_network_and_traj):
-    input_file = paths.Storage("setup.py", mode='w')
+    input_file = paths.Storage("setup.nc", mode='w')
     for obj in tps_network_and_traj:
         input_file.save(obj)
 
     input_file.tags['template'] = input_file.snapshots[0]
     input_file.close()
-    return "setup.py"
+    return "setup.nc"
 
 def test_append(tps_network_and_traj):
     runner = CliRunner()
     with runner.isolated_filesystem():
         in_file = make_input_file(tps_network_and_traj)
         result = runner.invoke(append, [in_file, '-a', 'output.nc',
                                         '--volume', 'A', '--volume', 'B'])
-        assert result.exit_code == 0
         assert result.exception is None
+        assert result.exit_code == 0
         storage = paths.Storage('output.nc', mode='r')
         assert len(storage.volumes) == 2
         assert len(storage.snapshots) == 0

diff --git a/paths_cli/tests/test_parameters.py b/paths_cli/tests/test_parameters.py
@@ -2,10 +2,44 @@
 import tempfile
 import os
 
-import openpathsampling as paths
+import paths_cli
 from openpathsampling.tests.test_helpers import make_1d_traj
 
 from paths_cli.parameters import *
+import openpathsampling as paths
+
+
+def pre_monkey_patch():
+    # store things that get monkey-patched; ensure we un-patch
+    stored_functions = {}
+    CallableCV = paths.CallableCV
+    PseudoAttr = paths.netcdfplus.FunctionPseudoAttribute
+    stored_functions['CallableCV.from'] = CallableCV.from_dict
+    stored_functions['PseudoAttr.from'] = PseudoAttr.from_dict
+    stored_functions['TPSNetwork.from'] = paths.TPSNetwork.from_dict
+    stored_functions['MISTISNetwork.from'] = paths.MISTISNetwork.from_dict
+    stored_functions['PseudoAttr.to'] = PseudoAttr.to_dict
+    stored_functions['TPSNetwork.to'] = paths.TPSNetwork.to_dict
+    stored_functions['MISTISNetwork.to'] = paths.MISTISNetwork.to_dict
+    return stored_functions
+
+def undo_monkey_patch(stored_functions):
+    CallableCV = paths.CallableCV
+    PseudoAttr = paths.netcdfplus.FunctionPseudoAttribute
+    CallableCV.from_dict = stored_functions['CallableCV.from']
+    PseudoAttr.from_dict = stored_functions['PseudoAttr.from']
+    paths.TPSNetwork.from_dict = stored_functions['TPSNetwork.from']
+    paths.MISTISNetwork.from_dict = stored_functions['MISTISNetwork.from']
+    PseudoAttr.to_dict = stored_functions['PseudoAttr.to']
+    paths.TPSNetwork.to_dict = stored_functions['TPSNetwork.to']
+    paths.MISTISNetwork.to_dict = stored_functions['MISTISNetwork.to']
+    paths_cli.param_core.StorageLoader.has_simstore_patch = False
+    paths.InterfaceSet.simstore = False
+    import importlib
+    importlib.reload(paths.netcdfplus)
+    importlib.reload(paths.collectivevariable)
+    importlib.reload(paths)
+
 
 
 class ParameterTest(object):
@@ -114,6 +148,17 @@ def setup(self):
     def test_get(self, getter):
         self._getter_test(getter)
 
+    def test_cannot_guess(self):
+        filename = self._filename('no-guess')
+        storage = paths.Storage(filename, 'w')
+        storage.save(self.engine)
+        storage.save(self.other_engine.named('other'))
+        storage.close()
+
+        storage = paths.Storage(filename, mode='r')
+        with pytest.raises(RuntimeError):
+            self.PARAMETER.get(storage, None)
+
 
 class TestSCHEME(ParamInstanceTest):
     PARAMETER = SCHEME
@@ -241,6 +286,17 @@ def test_get_multiple(self):
         assert traj0 == self.traj
         assert traj1 == self.other_traj
 
+    def test_cannot_guess(self):
+        filename = self._filename('no-guess')
+        storage = paths.Storage(filename, 'w')
+        storage.save(self.traj)
+        storage.save(self.other_traj)
+        storage.close()
+
+        storage = paths.Storage(filename, 'r')
+        with pytest.raises(RuntimeError):
+            self.PARAMETER.get(storage, None)
+
 
 class TestINIT_SNAP(ParamInstanceTest):
     PARAMETER = INIT_SNAP
@@ -278,6 +334,18 @@ def test_get(self, getter):
         obj = self.PARAMETER.get(storage, get_arg)
         assert obj == expected
 
+    def test_simstore_single_snapshot(self):
+        stored_functions = pre_monkey_patch()
+        filename = os.path.join(self.tempdir, "simstore.db")
+        storage = APPEND_FILE.get(filename)
+        storage.save(self.init_snap)
+        storage.close()
+
+        storage = INPUT_FILE.get(filename)
+        snap = self.PARAMETER.get(storage, None)
+        assert snap == self.init_snap
+        undo_monkey_patch(stored_functions)
+
 
 class MultiParamInstanceTest(ParamInstanceTest):
     def _getter_test(self, getter):
@@ -384,21 +452,26 @@ def test_get(self, getter):
         self._getter_test(getter)
 
 
-def test_OUTPUT_FILE():
+@pytest.mark.parametrize('ext', ['nc', 'db', 'sql'])
+def test_OUTPUT_FILE(ext):
+    stored_functions = pre_monkey_patch()
     tempdir = tempfile.mkdtemp()
-    filename = os.path.join(tempdir, "test_output_file.nc")
+    filename = os.path.join(tempdir, "test_output_file." + ext)
     assert not os.path.exists(filename)
     storage = OUTPUT_FILE.get(filename)
     assert os.path.exists(filename)
     os.remove(filename)
     os.rmdir(tempdir)
+    undo_monkey_patch(stored_functions)
 
-def test_APPEND_FILE():
+@pytest.mark.parametrize('ext', ['nc', 'db', 'sql'])
+def test_APPEND_FILE(ext):
+    stored_functions = pre_monkey_patch()
     tempdir = tempfile.mkdtemp()
-    filename = os.path.join(tempdir, "test_append_file.nc")
+    filename = os.path.join(tempdir, "test_append_file." + ext)
     assert not os.path.exists(filename)
     storage = APPEND_FILE.get(filename)
-    print(storage)
+    # print(storage)  # potentially useful debug; keep
     assert os.path.exists(filename)
     traj = make_1d_traj([0.0, 1.0])
     storage.tags['first_save'] = traj[0]
@@ -412,3 +485,4 @@ def test_APPEND_FILE():
     storage.close()
     os.remove(filename)
     os.rmdir(tempdir)
+    undo_monkey_patch(stored_functions)