Script for making zarr nodes

scripts/mknode.py

@@ -0,0 +1,297 @@
+#!/usr/bin/env -S uv run --script
+# /// script
+# requires-python = ">= 3.12"
+# dependencies = ["zarr >= 3.3.0"]
+# ///
+"""Script to create a new zarr node."""
+
+from __future__ import annotations
+
+import json
+import logging
+import sys
+from argparse import ArgumentParser
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from shutil import rmtree
+from typing import TypeVar
+
+import zarr
+from zarr.core.common import JSON, ZarrFormat
+
+logger = logging.getLogger("mknode")
+
+T = TypeVar("T")
+JSONObject = dict[str, JSON]
+
+ROOT_DIR_EXT = ".zarr"
+METADATA_FILE = "zarr.json"
+
+DATA_TYPES = ["bool"]
+for base in ("int", "uint"):
+    for precision in (8, 16, 32, 64):
+        DATA_TYPES.append(f"{base}{precision}")
+for precision in (32, 64):
+    DATA_TYPES.append(f"float{precision}")
+for precision in (64, 128):
+    DATA_TYPES.append(f"complex{precision}")
+
+
+def parse_list(string: str, fn: Callable[[str], T], sep: str = ",") -> list[T]:
+    return [fn(s.strip()) for s in string.split(sep)]
+
+
+def list_parser(fn: Callable[[str], T], sep: str = ",") -> Callable[[str], list[T]]:
+    return partial(parse_list, fn=fn, sep=sep)
+
+
+def jso(s: str) -> JSONObject:
+    j = json.loads(s)
+    if not isinstance(j, dict):
+        raise TypeError(f"Expected JSON object, got {s}")
+    return j
+
+
+@dataclass
+class ArrayArgs:
+    shape: list[int]
+    data_type: str
+    fill_value: JSON
+
+    @property
+    def chunk_shape(self) -> list[int]:
+        return self.shape.copy()
+
+    @classmethod
+    def maybe_from_args(
+        cls,
+        shape: list[int] | None,
+        data_type: str | None,
+        fill_value: JSON | None = None,
+    ):
+        if shape is None and data_type is None:
+            return None
+        if (shape is None) != (data_type is None):
+            raise ValueError("All array args must be given or none")
+        if fill_value is None:
+            if data_type == "bool":
+                fill_value = False
+            else:
+                fill_value = 0
+        return cls(shape, data_type, fill_value)  # type:ignore
+
+    def get_metadata(self, attributes: JSONObject | None = None) -> JSONObject:
+        if attributes is None:
+            attributes = {}
+
+        a2b: JSONObject = {"name": "bytes"}
+        if self.data_type not in ("bool", "int8", "uint8"):
+            a2b["configuration"] = {"endian": "little"}
+
+        d = {
+            "zarr_format": 3,
+            "node_type": "array",
+            "shape": self.shape,
+            "data_type": self.data_type,
+            "chunk_grid": {
+                "name": "regular",
+                "configuration": {"chunk_shape": self.shape},
+            },
+            "chunk_key_encoding": {"name": "default"},
+            "fill_value": self.fill_value,
+            "codecs": [a2b],
+            "attributes": attributes,
+        }
+        return d
+
+
+@dataclass
+class Args:
+    path: Path
+    store: Path
+    attributes: JSONObject
+    force: bool
+    parents: bool
+    log_level: int
+    zarr_version: ZarrFormat
+    array_args: ArrayArgs | None
+
+    def __post_init__(self):
+        if self.store != self.path and self.store not in self.path.parents:
+            raise ValueError("store must be an ancestor of path")
+
+    @classmethod
+    def parse(cls, raw_args: list[str] | None = None):
+        parser = ArgumentParser(description=__doc__)
+        parser.add_argument("path", type=Path, help="file system path to new node")
+        parser.add_argument(
+            "--zarr-version",
+            "-z",
+            choices=[2, 3],
+            type=int,
+            default=3,
+            help="Zarr version to write; default 3",
+        )
+        parser.add_argument(
+            "--store",
+            "-s",
+            type=Path,
+            help=f"file system path to store root, which must be an ancestor of the `path` argument; if not given, defaults to the nearest ancestor with the extension {ROOT_DIR_EXT}",
+        )
+        parser.add_argument(
+            "-a",
+            "--attributes",
+            type=jso,
+            help="attributes to add to the new node, as a JSON string representing an object",
+        )
+        parser.add_argument(
+            "-f",
+            "--force",
+            action="store_true",
+            default=False,
+            help="if the node already exists, delete it",
+        )
+        parser.add_argument(
+            "-p",
+            "--parents",
+            action="store_true",
+            default=False,
+            help="create parent Zarr groups, including the store root and its parent directory, if necessary",
+        )
+        parser.add_argument(
+            "-v",
+            "--verbose",
+            action="count",
+            default=0,
+            help="increase logging verbosity",
+        )
+        g = parser.add_argument_group(
+            "array", "Additional arguments for creating an array rather than a group."
+        )
+        g.add_argument(
+            "shape",
+            type=list_parser(int),
+            nargs="?",
+            help="comma-separated list of unsigned integers representing array shape",
+        )
+        g.add_argument(
+            "datatype",
+            nargs="?",
+            choices=DATA_TYPES,
+            help="data type for the array",
+        )
+        g.add_argument(
+            "--fill-value",
+            "-F",
+            type=jso,
+            help="fill value to be used as JSON; not type-checked",
+        )
+        parsed = parser.parse_args(raw_args)
+        maybe_array = ArrayArgs.maybe_from_args(
+            parsed.shape, parsed.datatype, parsed.fill_value
+        )
+        level = {0: logging.WARNING, 1: logging.INFO, 2: logging.DEBUG}.get(
+            parsed.verbose, logging.DEBUG
+        )
+        storepath: Path | None = parsed.store
+        nodepath: Path = parsed.path
+        if storepath is None:
+            if nodepath.name.endswith(ROOT_DIR_EXT):
+                storepath = nodepath
+            else:
+                for p in nodepath.parents:
+                    if p.name.endswith(ROOT_DIR_EXT):
+                        logger.info("Inferring %s as store root", p)
+                        storepath = p
+                        break
+
+        if storepath is None:
+            logger.warning(
+                "No --store given, and could not infer from %s extension; node path will be used as store root, but should be renamed",
+                ROOT_DIR_EXT,
+            )
+            storepath = nodepath
+
+        return cls(
+            nodepath,
+            storepath,
+            parsed.attributes or {},
+            parsed.force,
+            parsed.parents,
+            level,
+            parsed.zarr_version,
+            maybe_array,
+        )
+
+
+def eprint(*args, **kwargs):
+    kwargs.setdefault("file", sys.stderr)
+    print(*args, **kwargs)
+
+
+def grp_metadata(attrs: JSONObject | None = None) -> JSONObject:
+    if attrs is None:
+        attrs = {}
+    return {"zarr_format": 3, "node_type": "group", "attributes": attrs}
+
+
+def write_node_metadata(path: Path, metadata: JSONObject):
+    s = json.dumps(metadata, indent=2, sort_keys=True) + "\n"
+    p = path.joinpath(METADATA_FILE)
+    p.write_text(s)
+    if logger.isEnabledFor(logging.INFO):
+        logger.info(
+            "Wrote metadata into %s : %s", p, json.dumps(metadata, sort_keys=True)
+        )
+
+
+def write_group_metadata(path: Path, attrs: JSONObject | None = None):
+    write_node_metadata(path, grp_metadata(attrs))
+
+
+def main():
+    args = Args.parse()
+    logging.basicConfig(level=args.log_level)
+
+    if args.store != args.path:
+        args.store.parent.mkdir(exist_ok=True, parents=args.parents)
+        mode = "a" if args.parents else "r+"
+        grp = zarr.open_group(args.store, mode=mode, zarr_format=args.zarr_version)
+        for name in args.path.relative_to(args.store).parts[:-1]:
+            if args.parents:
+                grp = grp.require_group(name)
+            else:
+                grp = grp.get_group(name)
+
+    nodepath = args.path
+    if nodepath.exists():
+        if args.force:
+            logger.warning("Removing existing node at %s", nodepath)
+            rmtree(nodepath)
+        else:
+            eprint(f"Node already exists at {nodepath} ; use --force to overwrite")
+            return 1
+
+    if args.array_args is None:
+        zarr.create_group(
+            args.path, zarr_format=args.zarr_version, attributes=args.attributes
+        )
+    else:
+        aargs = args.array_args
+        zarr.create_array(
+            args.path,
+            shape=aargs.shape,
+            dtype=aargs.data_type,
+            chunks=aargs.chunk_shape,
+            fill_value=aargs.fill_value,
+            attributes=args.attributes,
+        )
+
+    return 0
+
+
+if __name__ == "__main__":
+    status = main()
+    sys.exit(status)