Merge pull request #2 from dwhswenson/append

Appending to files

Author: dwhswenson

paths_cli/cli.py

@@ -7,6 +7,8 @@
 import os
 
 import click
+# import click_completion
+# click_completion.init()
 
 CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])
 
@@ -116,14 +118,14 @@ def format_commands(self, ctx, formatter):
     openpathsampling strip-snapshots --help
 """
 
+OPS_CLI = OpenPathSamplingCLI(
+    name="openpathsampling",
+    help=_MAIN_HELP,
+    context_settings=CONTEXT_SETTINGS
+)
 
 def main():  # no-cov
-    cli = OpenPathSamplingCLI(
-        name="openpathsampling",
-        help=_MAIN_HELP,
-        context_settings=CONTEXT_SETTINGS
-    )
-    cli()
+    OPS_CLI()
 
 
 if __name__ == '__main__':  # no-cov

paths_cli/commands/append.py

@@ -0,0 +1,59 @@
+import click
+from paths_cli.parameters import (
+    INPUT_FILE, APPEND_FILE, MULTI_CV, MULTI_ENGINE, MULTI_VOLUME,
+    MULTI_NETWORK, MULTI_SCHEME, MULTI_TAG
+)
+
+@click.command(
+    'append',
+    short_help="add objects from INPUT_FILE  to another file"
+)
+@INPUT_FILE.clicked(required=True)
+@APPEND_FILE.clicked(required=True)
+@MULTI_ENGINE.clicked(required=False)
+@MULTI_CV.clicked(required=False)
+@MULTI_VOLUME.clicked(required=False)
+@MULTI_NETWORK.clicked(required=False)
+@MULTI_SCHEME.clicked(required=False)
+@MULTI_TAG.clicked(required=False)
+@click.option('--save-tag', type=str, default=None,
+              help=("save object to a tag; requires that only one "
+                    + "object be specfied. Can also be used to rename "
+                    + "tagged objects. To append a tagged object without "
+                    + "a tag, use --save-tag \"\""))
+def append(input_file, append_file, engine, cv, volume, network, scheme,
+           tag, save_tag):
+    """Append objects from INPUT_FILE to another file.
+    """
+    storage = INPUT_FILE.get(input_file)
+    output_storage = APPEND_FILE.get(append_file)
+    params = [MULTI_ENGINE, MULTI_CV, MULTI_VOLUME, MULTI_NETWORK,
+              MULTI_SCHEME, MULTI_TAG]
+    args = [engine, cv, volume, network, scheme, tag]
+    to_save = []
+    for arg, param in zip(args, params):
+        to_save.extend(param.get(storage, arg))
+
+    to_save = [obj for obj in to_save if obj is not None]
+    if save_tag is not None and len(to_save) != 1:
+        raise RuntimeError("Can't identify the object to tag when saving "
+                           + str(len(to_save)) + " objects.")
+
+    for obj in to_save:
+        output_storage.save(obj)
+
+    if tag and len(tag) == 1 and save_tag is None:
+        save_tag = tag[0]
+
+    if save_tag:
+        output_storage.tags[save_tag] = to_save[0]
+
+    # TO TEST
+    # 3. "untag" an object by not associating a tag in the new storage
+
+    output_storage.close()
+
+
+CLI = append
+SECTION = "Miscellaneous"
+REQUIRES_OPS = (1, 0)

paths_cli/parameters.py

@@ -1,4 +1,5 @@
 import click
+import os
 # import openpathsampling as paths
 
 class AbstractParameter(object):
@@ -14,6 +15,8 @@ def clicked(self, required=False):
 # we'll use tests of the -h option in the .travis.yml to ensure that the
 # .clicked methods work
 
+HELP_MULTIPLE = "; may be used more than once"
+
 class Option(AbstractParameter):
     def clicked(self, required=False):  # no-cov
         return click.option(*self.args, **self.kwargs, required=required)
@@ -40,18 +43,38 @@ def __init__(self, param, mode):
         super(StorageLoader, self).__init__(param)
         self.mode = mode
 
+    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):
+            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)
 
 
 class OPSStorageLoadNames(AbstractLoader):
+    """Simple loader that expects its input to be a name or index.
+    """
     def __init__(self, param, store):
         super(OPSStorageLoadNames, self).__init__(param)
         self.store = store
 
     def get(self, storage, names):
-        return [getattr(storage, self.store)[name] for name in names]
+        int_corrected = []
+        for name in names:
+            try:
+                name = int(name)
+            except ValueError:
+                pass
+            int_corrected.append(name)
+
+        return [getattr(storage, self.store)[name]
+                for name in int_corrected]
 
 
 class OPSStorageLoadSingle(AbstractLoader):
@@ -171,16 +194,47 @@ def init_snap_fallback(parameter, storage, name):
 
 CVS = OPSStorageLoadNames(
     param=Option('--cv', type=str, multiple=True,
-                 help='name of CV; may select more than once'),
+                 help='name of CV' + HELP_MULTIPLE),
     store='cvs'
 )
 
+MULTI_VOLUME = OPSStorageLoadNames(
+    param=Option('--volume', type=str, multiple=True,
+                 help='name or index of volume' + HELP_MULTIPLE),
+    store='volumes'
+)
+
+MULTI_ENGINE = OPSStorageLoadNames(
+    param=Option('--engine', type=str, multiple=True,
+                 help='name or index of engine' + HELP_MULTIPLE),
+    store='engines'
+)
+
+
 STATES = OPSStorageLoadNames(
-    param=Option('-s', '--state', multiple=True,
-                 help='name of state; may select more than once'),
+    param=Option('-s', '--state', type=str, multiple=True,
+                 help='name  of state' + HELP_MULTIPLE),
     store='volumes'
 )
 
+MULTI_TAG = OPSStorageLoadNames(
+    param=Option('--tag', type=str, multiple=True,
+                 help='tag for object' + HELP_MULTIPLE),
+    store='tags'
+)
+
+MULTI_NETWORK = OPSStorageLoadNames(
+    param=Option('--network', type=str, multiple=True,
+                 help='name or index of network' + HELP_MULTIPLE),
+    store='networks'
+)
+
+MULTI_SCHEME = OPSStorageLoadNames(
+    param=Option('--scheme', type=str, multiple=True,
+                 help='name or index of move scheme' + HELP_MULTIPLE),
+    store='schemes'
+)
+
 INPUT_FILE = StorageLoader(
     param=Argument('input_file',
                    type=click.Path(exists=True, readable=True)),
@@ -194,5 +248,14 @@ def init_snap_fallback(parameter, storage, name):
     mode='w'
 )
 
+APPEND_FILE = StorageLoader(
+    param=Option('-a', '--append-file',
+                 type=click.Path(writable=True, readable=True),
+                 help="file to append to"),
+    mode='a'
+)
+
 N_STEPS_MC = click.option('-n', '--nsteps', type=int,
                           help="number of Monte Carlo trials to run")
+
+MULTI_CV = CVS

paths_cli/tests/commands/test_append.py

@@ -0,0 +1,99 @@
+import os
+
+import pytest
+from click.testing import CliRunner
+
+from paths_cli.commands.append import *
+
+import openpathsampling as paths
+
+def make_input_file(tps_network_and_traj):
+    input_file = paths.Storage("setup.py", 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"
+
+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
+        storage = paths.Storage('output.nc', mode='r')
+        assert len(storage.volumes) == 2
+        assert len(storage.snapshots) == 0
+        storage.volumes['A']  # smoke tests that we can load
+        storage.volumes['B']
+        storage.close()
+
+        result = runner.invoke(append, [in_file, '-a', 'output.nc',
+                                        '--tag', 'template'])
+        storage = paths.Storage('output.nc', mode='r')
+        assert len(storage.volumes) == 2
+        assert len(storage.snapshots) == 2  # one snapshot + reverse
+
+@pytest.mark.parametrize('n_objects', [0, 2])
+def test_append_tag_error(tps_network_and_traj, n_objects):
+    objs = {2: ['--volume', "A", '--volume', "B"], 0: []}[n_objects]
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        in_file = make_input_file(tps_network_and_traj)
+        result = runner.invoke(append,
+                               [in_file, '-a', "output.nc"] + objs
+                               + ["--save-tag", "foo"])
+        assert isinstance(result.exception, RuntimeError)
+        assert "Can't identify the object to tag" in str(result.exception)
+
+def test_append_tag(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",
+                                '--tag', 'template', '--save-tag', 'foo'])
+        assert result.exit_code == 0
+        assert result.exception is None
+
+        storage = paths.Storage("output.nc", mode='r')
+        assert len(storage.snapshots) == 2
+        assert len(storage.tags) == 1
+        assert storage.tags['foo'] is not None
+        storage.close()
+
+def test_append_same_tag(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",
+                                '--tag', 'template'])
+        assert result.exit_code == 0
+        assert result.exception is None
+
+        storage = paths.Storage("output.nc", mode='r')
+        assert len(storage.snapshots) == 2
+        assert len(storage.tags) == 1
+        assert storage.tags['template'] is not None
+        storage.close()
+
+def test_append_remove_tag(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",
+                                "--tag", 'template', '--save-tag', ''])
+        print(result.output)
+        assert result.exception is None
+        assert result.exit_code == 0
+
+        storage = paths.Storage("output.nc", mode='r')
+        assert len(storage.snapshots) == 2
+        assert len(storage.tags) == 0
+        storage.close()
+