Repository: isi-vista/vistautils

Issue:
Script to turn a directory into a key-value zip usable as a Key-Value Source
We often need to pass experimental artifacts from one user to another.  Currently we mostly do this using lists of file paths together with the files themselves, but this is inconvenient when the data moves around (e.g. from a local laptop to our network code).

We have `KeyValueSource`s and `KeyValueSink`s backed by zip files which work better for this purpose. However, if we have incoming data which is not already stored as one of these, it would be helpful to have a script to convert a list of files in a directory into one. (Just using `zip` won't write the proper indices).

Comments:

Patch:
diff --git a/vistautils/scripts/directory_to_key_value_store.py b/vistautils/scripts/directory_to_key_value_store.py
new file mode 100644
index 0000000..bc2855b
--- /dev/null
+++ b/vistautils/scripts/directory_to_key_value_store.py
@@ -0,0 +1,74 @@
+"""
+Given a directory, copies its content into an object which can be used as a `KeyValueSource`.
+
+The form of the output key-value store is given by the "output" sub-namespace of the parameters.
+See `byte_key_value_sink_from_params` for details.
+In the most common case of a zip-backed store, create a sub-namespace like this:
+.. code-block:: YAML
+
+    output:
+      type: zip
+      path: path of zip to write to
+
+Each file in the input directory will become one entry in the output store
+with a value equal to the binary content of the file.
+The key will by default be the file name (not the full path).
+This can be changed by specifying the *key_function* parameter.
+Currently the only key function supported is *strip_one_extension*,
+which removes the first extension (if present) from the file name before using it as the key.
+
+Because the correct way to handle them is unclear,
+if any sub-directories of the input directory are encountered, an error will be raised.
+This behavior is subject to change in the future and should not be relied upon.
+"""
+import logging
+from pathlib import Path
+from typing import Callable
+
+from vistautils.key_value import byte_key_value_sink_from_params
+from vistautils.parameters import Parameters
+from vistautils.parameters_only_entrypoint import parameters_only_entry_point
+
+
+def main(params: Parameters) -> None:
+    input_directory = params.existing_directory("input_directory")
+    key_function = key_function_from_params(params)
+
+    with byte_key_value_sink_from_params(params, eval_context=locals()) as sink:
+        for item_path in input_directory.glob("*"):
+            if item_path.is_dir():
+                raise NotImplementedError(
+                    f"We don't know how to handle sub-directories in the input, but got "
+                    f"{item_path}"
+                )
+            logging.info("Copying %s to output sink", item_path)
+            sink.put(key=key_function(item_path), value=item_path.read_bytes())
+
+
+def key_function_from_params(params: Parameters) -> Callable[[Path], str]:
+    key_function_string = (
+        params.optional_string("key_function", [IDENTITY, STRIP_ONE_EXTENSION])
+        or IDENTITY
+    )
+    if key_function_string == IDENTITY:
+        return identity_key_function
+    elif key_function_string == STRIP_ONE_EXTENSION:
+        return strip_one_extension_key_function
+    else:
+        raise NotImplementedError(f"Unknown key function %s", key_function_string)
+
+
+IDENTITY = "identity"
+STRIP_ONE_EXTENSION = "strip_one_extension"
+
+
+def identity_key_function(path: Path) -> str:
+    return path.name
+
+
+def strip_one_extension_key_function(path: Path) -> str:
+    return path.stem
+
+
+if __name__ == "__main__":
+    parameters_only_entry_point(main)


Test patch:
diff --git a/tests/scripts/__init__.py b/tests/scripts/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/scripts/test_directory_to_key_value_store.py b/tests/scripts/test_directory_to_key_value_store.py
new file mode 100644
index 0000000..2d2abd3
--- /dev/null
+++ b/tests/scripts/test_directory_to_key_value_store.py
@@ -0,0 +1,43 @@
+from pathlib import Path
+from typing import Optional
+
+from vistautils.key_value import KeyValueSource
+from vistautils.parameters import Parameters
+from vistautils.scripts import directory_to_key_value_store
+
+import pytest
+
+
+@pytest.mark.parametrize("key_function", [None, "strip_one_extension"])
+def test_directory_key_value_store(tmp_path: Path, key_function: Optional[str]):
+    input_dir = tmp_path / "input_dir"
+    input_dir.mkdir(exist_ok=True)
+
+    # We test filenames with 0, 1, and 2 extensions.
+    filenames = ("fred", "bob.txt", "melvin.tar.gz")
+    filenames_without_one_extension = ("fred", "bob", "melvin.tar")
+    file_content = ("hello", "world", "foo")
+
+    for (filename, content) in zip(filenames, file_content):
+        (input_dir / filename).write_text(content, encoding="utf-8")
+
+    output_zip_path = (tmp_path / "output.zip").absolute()
+    params = {
+        "input_directory": str(input_dir.absolute()),
+        "output": {"type": "zip", "path": str(output_zip_path)},
+    }
+
+    if key_function:
+        params["key_function"] = key_function
+
+    directory_to_key_value_store.main(Parameters.from_mapping(params))
+
+    if key_function == "strip_one_extension":
+        reference = dict(zip(filenames_without_one_extension, file_content))
+    else:
+        reference = dict(zip(filenames, file_content))
+
+    with KeyValueSource.zip_character_source(output_zip_path) as source:
+        assert set(source.keys()) == set(reference.keys())
+        for key, reference_value in reference.items():
+            assert source[key] == reference_value