diff --git a/pyproject.toml b/pyproject.toml index 07664b095..f8b8c88d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,7 @@ fsspec = ["fsspec[http]"] [project.scripts] parquet-to-blosc2 = "blosc2.cli.parquet_to_blosc2:main" b2nd-to-zarr = "blosc2.cli.b2nd_to_zarr:main" +zarr-to-b2nd = "blosc2.cli.zarr_to_b2nd:main" blosc2-to-zarr = "blosc2.cli.b2nd_to_zarr:main" b2view = "blosc2.b2view.cli:main" diff --git a/src/blosc2/cli/b2nd_to_zarr.py b/src/blosc2/cli/b2nd_to_zarr.py index abad6745e..17dd024bb 100644 --- a/src/blosc2/cli/b2nd_to_zarr.py +++ b/src/blosc2/cli/b2nd_to_zarr.py @@ -11,126 +11,25 @@ import argparse import contextlib -import itertools import math -import os import shutil import sys import time from pathlib import Path from typing import Any -import numpy as np - import blosc2 - -try: - from rich.console import Console - from rich.progress import ( - BarColumn, - Progress, - TaskProgressColumn, - TextColumn, - TimeElapsedColumn, - TimeRemainingColumn, - TransferSpeedColumn, - ) - from rich.table import Table - - HAVE_RICH = True -except ImportError: - HAVE_RICH = False - -DEFAULT_BUFFER_SIZE_MB = 128 - - -def require_zarr() -> Any: - """Ensure the zarr package is installed.""" - try: - import zarr - - return zarr - except ImportError as exc: - raise ImportError("b2nd-to-zarr requires zarr; install it with: pip install 'blosc2[zarr]'") from exc - - -def format_bytes(size: float) -> str: - """Format byte count into human-readable string.""" - units = ["B", "KB", "MB", "GB", "TB"] - unit_idx = 0 - while size >= 1024.0 and unit_idx < len(units) - 1: - size /= 1024.0 - unit_idx += 1 - return f"{size:,.2f} {units[unit_idx]}" if unit_idx > 0 else f"{int(size)} B" - - -def get_path_size(path: Path) -> int: - """Get total byte size of a file or directory tree.""" - if not path.exists(): - return 0 - if path.is_file(): - return path.stat().st_size - total = 0 - for root, _, files in os.walk(path): - for f in files: - fp = os.path.join(root, f) - with contextlib.suppress(OSError): - total += os.path.getsize(fp) - return total - - -def parse_shape(shape_str: str) -> tuple[int, ...]: - """Parse comma- or whitespace-separated integer dimensions.""" - cleaned = shape_str.replace("x", ",").replace(" ", ",").replace("(", "").replace(")", "") - dims = [int(p.strip()) for p in cleaned.split(",") if p.strip()] - if not dims or any(d <= 0 for d in dims): - raise ValueError(f"Invalid shape string '{shape_str}', dimensions must be positive integers.") - return tuple(dims) - - -def compute_copy_shape( - shape: tuple[int, ...], - target_chunks: tuple[int, ...], - itemsize: int, - max_buffer_bytes: int = DEFAULT_BUFFER_SIZE_MB * 1024 * 1024, -) -> tuple[int, ...]: - """Compute batch copy shape aligned with chunk boundaries to optimize memory and throughput.""" - copy_shape = list(target_chunks) - ndim = len(shape) - - # First attempt: expand inner dimensions to full shape (from right to left) - # to maximize contiguous memory slicing in C-order arrays. - for d in reversed(range(ndim)): - test_shape = list(copy_shape) - test_shape[d] = shape[d] - test_bytes = math.prod(test_shape) * itemsize - if test_bytes <= max_buffer_bytes: - copy_shape[d] = shape[d] - else: - current_bytes = math.prod(copy_shape) * itemsize - if current_bytes < max_buffer_bytes: - bytes_per_chunk = current_bytes // math.ceil(copy_shape[d] / target_chunks[d]) - max_chunks = max_buffer_bytes // bytes_per_chunk - num_chunks = math.ceil(shape[d] / target_chunks[d]) - mult = min(max(1, max_chunks), num_chunks) - copy_shape[d] = min(shape[d], mult * target_chunks[d]) - break - - # Second attempt: if memory headroom remains, expand leading dimensions - for d in range(ndim): - current_bytes = math.prod(copy_shape) * itemsize - if current_bytes >= max_buffer_bytes: - break - num_chunks = math.ceil(shape[d] / target_chunks[d]) - curr_chunks = math.ceil(copy_shape[d] / target_chunks[d]) - if curr_chunks < num_chunks: - bytes_per_chunk = current_bytes // curr_chunks - max_chunks = max_buffer_bytes // bytes_per_chunk - mult = min(max(1, max_chunks), num_chunks) - if mult > curr_chunks: - copy_shape[d] = min(shape[d], mult * target_chunks[d]) - - return tuple(copy_shape) +from blosc2.cli.conversion_utils import ( + DEFAULT_BUFFER_SIZE_MB, + compute_batches, + copy_batches, + format_bytes, + get_path_size, + parse_shape, + print_summary, + require_zarr, + verify_arrays, +) def resolve_shuffle(arr: blosc2.NDArray, shuffle_mode: str) -> str: @@ -251,113 +150,6 @@ def copy_metadata(arr: blosc2.NDArray, z_arr: Any) -> None: z_arr.attrs[k] = v -def copy_batches( - arr: blosc2.NDArray, - z_arr: Any, - batch_slices: list[tuple[slice, ...]], - total_bytes: int, - quiet: bool = False, - verbose: bool = False, -) -> float: - """Copy data slices from Blosc2 NDArray to Zarr array with progress tracking.""" - total_batches = len(batch_slices) - if total_batches == 0: - return 0.0 - t_start = time.perf_counter() - - if HAVE_RICH and not quiet: - console = Console() - with Progress( - TextColumn("[bold blue]{task.description}"), - BarColumn(), - TaskProgressColumn(), - TransferSpeedColumn(), - TimeElapsedColumn(), - TimeRemainingColumn(), - console=console, - ) as progress: - task = progress.add_task("Converting b2nd -> zarr...", total=total_bytes) - for s in batch_slices: - data = arr[s] - z_arr[s] = data - batch_size = data.nbytes - progress.update(task, advance=batch_size) - del data - else: - for idx, s in enumerate(batch_slices): - data = arr[s] - z_arr[s] = data - del data - if verbose and not quiet: - pct = (idx + 1) / total_batches * 100 - print(f"Batch {idx + 1}/{total_batches} ({pct:.1f}%) written") - - return time.perf_counter() - t_start - - -def _fast_sample_verify( - b2_arr: blosc2.NDArray, z_arr: Any, slices_list: list[tuple[slice, ...]], quiet: bool -) -> None: - corner_indices = [ - tuple(0 for _ in b2_arr.shape), - tuple(-1 for _ in b2_arr.shape), - tuple(s // 2 for s in b2_arr.shape), - ] - for idx in corner_indices: - v_b2 = b2_arr[idx] - v_z = z_arr[idx] - if not np.array_equal(v_b2, v_z, equal_nan=True): - raise ValueError(f"Verification failed at index {idx}: Blosc2={v_b2} != Zarr={v_z}") - - if slices_list: - for s in (slices_list[0], slices_list[-1]): - b2_chunk = b2_arr[s] - z_chunk = z_arr[s] - if not np.array_equal(b2_chunk, z_chunk, equal_nan=True): - raise ValueError(f"Verification failed for slice {s}") - if not quiet: - print("✓ Fast sample verification passed!") - - -def _full_verify( - b2_arr: blosc2.NDArray, z_arr: Any, slices_list: list[tuple[slice, ...]], quiet: bool -) -> None: - if not quiet: - print("Performing full verification across all slices...") - for idx, s in enumerate(slices_list): - b2_data = b2_arr[s] - z_data = z_arr[s] - if not np.array_equal(b2_data, z_data, equal_nan=True): - raise ValueError(f"Full verification failed at slice {s} (batch {idx})") - if not quiet: - print("✓ Full element verification passed!") - - -def verify_arrays( - b2_arr: blosc2.NDArray, - z_arr: Any, - slices_list: list[tuple[slice, ...]], - full: bool = False, - quiet: bool = False, -) -> bool: - """Verify that data in Zarr array matches source Blosc2 NDArray.""" - if b2_arr.shape != z_arr.shape: - raise ValueError(f"Shape mismatch: Blosc2 {b2_arr.shape} vs Zarr {z_arr.shape}") - if b2_arr.dtype != z_arr.dtype: - raise ValueError(f"Dtype mismatch: Blosc2 {b2_arr.dtype} vs Zarr {z_arr.dtype}") - - if math.prod(b2_arr.shape) == 0: - if not quiet: - print("✓ Empty array verification passed!") - return True - - if not full: - _fast_sample_verify(b2_arr, z_arr, slices_list, quiet) - else: - _full_verify(b2_arr, z_arr, slices_list, quiet) - return True - - def _determine_target_chunks_and_shards( arr: blosc2.NDArray, chunks: tuple[int, ...] | None, @@ -381,26 +173,6 @@ def _determine_target_chunks_and_shards( return target_chunks, target_shards -def _compute_batches( - shape: tuple[int, ...], - target_chunks: tuple[int, ...], - target_shards: tuple[int, ...] | None, - itemsize: int, - buffer_size_mb: int, -) -> tuple[tuple[int, ...], list[tuple[slice, ...]]]: - if math.prod(shape) == 0: - return shape, [] - unit_shape = target_shards if target_shards is not None else target_chunks - copy_shape = compute_copy_shape( - shape, unit_shape, itemsize, max_buffer_bytes=buffer_size_mb * 1024 * 1024 - ) - slices_per_dim = [ - [slice(i, min(i + step, s)) for i in range(0, s, step)] - for s, step in zip(shape, copy_shape, strict=True) - ] - return copy_shape, list(itertools.product(*slices_per_dim)) - - def b2nd_to_zarr( input_path: str | Path, output_path: str | Path | None = None, @@ -421,7 +193,7 @@ def b2nd_to_zarr( verbose: bool = False, ) -> dict[str, Any]: """Convert a Blosc2 .b2nd file containing an NDArray to a Zarr array.""" - zarr = require_zarr() + zarr = require_zarr("b2nd-to-zarr") src_path = Path(input_path).expanduser().resolve() if not src_path.exists(): @@ -465,9 +237,7 @@ def b2nd_to_zarr( blocksize=blocksize, zarr_format=zarr_format, ) - copy_shape, batch_slices = _compute_batches( - shape, target_chunks, target_shards, itemsize, buffer_size_mb - ) + copy_shape, batch_slices = compute_batches(shape, target_chunks, target_shards, itemsize, buffer_size_mb) if verbose and not quiet: print(f"Source: {src_path} ({format_bytes(src_file_size)})") @@ -504,7 +274,9 @@ def b2nd_to_zarr( ) copy_metadata(arr, z_arr) - copy_duration = copy_batches(arr, z_arr, batch_slices, uncompressed_bytes, quiet=quiet, verbose=verbose) + copy_duration = copy_batches( + arr, z_arr, batch_slices, uncompressed_bytes, "b2nd -> zarr", quiet=quiet, verbose=verbose + ) if zip_store is not None: zip_store.close() @@ -706,44 +478,7 @@ def main(argv: list[str] | None = None) -> int: return 1 if not args.quiet: - if HAVE_RICH: - console = Console() - table = Table(title="[bold green]Conversion Summary[/bold green]", box=None) - table.add_column("Property", style="bold cyan") - table.add_column("Value") - table.add_row("Source", f"{summary['src_path']} ({format_bytes(summary['src_size_bytes'])})") - table.add_row( - "Destination", f"{summary['dst_path']} ({format_bytes(summary['dst_size_bytes'])})" - ) - table.add_row("Shape / Dtype", f"{summary['shape']} / {summary['dtype']}") - table.add_row("Chunks", str(summary["chunks"])) - if summary["shards"] is not None: - table.add_row("Shards", str(summary["shards"])) - table.add_row("Uncompressed Size", format_bytes(summary["uncompressed_bytes"])) - table.add_row("Source Compression", f"{summary['src_cratio']:.2f}x") - table.add_row("Zarr Compression", f"{summary['dst_cratio']:.2f}x") - table.add_row( - "Conversion Speed", - f"{summary['throughput_mb_s']:.1f} MB/s ({summary['copy_duration_sec']:.2f} s)", - ) - table.add_row("Total Time", f"{summary['total_duration_sec']:.2f} s") - console.print() - console.print(table) - else: - print("\n--- Conversion Summary ---") - print(f"Source: {summary['src_path']} ({format_bytes(summary['src_size_bytes'])})") - print(f"Destination: {summary['dst_path']} ({format_bytes(summary['dst_size_bytes'])})") - print(f"Shape / Dtype: {summary['shape']} / {summary['dtype']}") - print(f"Chunks: {summary['chunks']}") - if summary["shards"] is not None: - print(f"Shards: {summary['shards']}") - print(f"Uncompressed: {format_bytes(summary['uncompressed_bytes'])}") - print(f"Source CRatio: {summary['src_cratio']:.2f}x") - print(f"Zarr CRatio: {summary['dst_cratio']:.2f}x") - print( - f"Conversion Speed: {summary['throughput_mb_s']:.1f} MB/s ({summary['copy_duration_sec']:.2f} s)" - ) - print(f"Total Time: {summary['total_duration_sec']:.2f} s") + print_summary(summary, dst_label="Zarr") return 0 diff --git a/src/blosc2/cli/conversion_utils.py b/src/blosc2/cli/conversion_utils.py new file mode 100644 index 000000000..3cde82dea --- /dev/null +++ b/src/blosc2/cli/conversion_utils.py @@ -0,0 +1,310 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +"""Conversion utilities for converting between a Blosc2 .b2nd file (NDArray) and a Zarr array.""" + +from __future__ import annotations + +import contextlib +import itertools +import math +import os +import time +from typing import TYPE_CHECKING, Any + +import numpy as np + +import blosc2 + +if TYPE_CHECKING: + from pathlib import Path + +try: + from rich.console import Console + from rich.progress import ( + BarColumn, + Progress, + TaskProgressColumn, + TextColumn, + TimeElapsedColumn, + TimeRemainingColumn, + TransferSpeedColumn, + ) + from rich.table import Table + + HAVE_RICH = True +except ImportError: + HAVE_RICH = False + +DEFAULT_BUFFER_SIZE_MB = 128 + +BLOSC2_CODECS = { + "blosclz": blosc2.Codec.BLOSCLZ, + "lz4": blosc2.Codec.LZ4, + "lz4hc": blosc2.Codec.LZ4HC, + "zlib": blosc2.Codec.ZLIB, + "zstd": blosc2.Codec.ZSTD, +} + +BLOSC2_SHUFFLE_FILTERS = { + "noshuffle": blosc2.Filter.NOFILTER, + "shuffle": blosc2.Filter.SHUFFLE, + "bitshuffle": blosc2.Filter.BITSHUFFLE, +} + + +def format_bytes(size: float) -> str: + """Format byte count into human-readable string.""" + units = ["B", "KB", "MB", "GB", "TB"] + unit_idx = 0 + while size >= 1024.0 and unit_idx < len(units) - 1: + size /= 1024.0 + unit_idx += 1 + return f"{size:,.2f} {units[unit_idx]}" if unit_idx > 0 else f"{int(size)} B" + + +def get_path_size(path: Path) -> int: + """Get total byte size of a file or directory tree.""" + if not path.exists(): + return 0 + if path.is_file(): + return path.stat().st_size + total = 0 + for root, _, files in os.walk(path): + for f in files: + fp = os.path.join(root, f) + with contextlib.suppress(OSError): + total += os.path.getsize(fp) + return total + + +def parse_shape(shape_str: str) -> tuple[int, ...]: + """Parse comma- or whitespace-separated integer dimensions.""" + cleaned = shape_str.replace("x", ",").replace(" ", ",").replace("(", "").replace(")", "") + dims = [int(p.strip()) for p in cleaned.split(",") if p.strip()] + if not dims or any(d <= 0 for d in dims): + raise ValueError(f"Invalid shape string '{shape_str}', dimensions must be positive integers.") + return tuple(dims) + + +def compute_copy_shape( + shape: tuple[int, ...], + target_chunks: tuple[int, ...], + itemsize: int, + max_buffer_bytes: int = DEFAULT_BUFFER_SIZE_MB * 1024 * 1024, +) -> tuple[int, ...]: + """Compute batch copy shape aligned with chunk boundaries to optimize memory and throughput.""" + copy_shape = list(target_chunks) + ndim = len(shape) + + # First attempt: expand inner dimensions to full shape (from right to left) + # to maximize contiguous memory slicing in C-order arrays. + for d in reversed(range(ndim)): + test_shape = list(copy_shape) + test_shape[d] = shape[d] + test_bytes = math.prod(test_shape) * itemsize + if test_bytes <= max_buffer_bytes: + copy_shape[d] = shape[d] + else: + current_bytes = math.prod(copy_shape) * itemsize + if current_bytes < max_buffer_bytes: + bytes_per_chunk = current_bytes // math.ceil(copy_shape[d] / target_chunks[d]) + max_chunks = max_buffer_bytes // bytes_per_chunk + num_chunks = math.ceil(shape[d] / target_chunks[d]) + mult = min(max(1, max_chunks), num_chunks) + copy_shape[d] = min(shape[d], mult * target_chunks[d]) + break + + # Second attempt: if memory headroom remains, expand leading dimensions + for d in range(ndim): + current_bytes = math.prod(copy_shape) * itemsize + if current_bytes >= max_buffer_bytes: + break + num_chunks = math.ceil(shape[d] / target_chunks[d]) + curr_chunks = math.ceil(copy_shape[d] / target_chunks[d]) + if curr_chunks < num_chunks: + bytes_per_chunk = current_bytes // curr_chunks + max_chunks = max_buffer_bytes // bytes_per_chunk + mult = min(max(1, max_chunks), num_chunks) + if mult > curr_chunks: + copy_shape[d] = min(shape[d], mult * target_chunks[d]) + + return tuple(copy_shape) + + +def compute_batches( + shape: tuple[int, ...], + target_chunks: tuple[int, ...], + target_shards: tuple[int, ...] | None, + itemsize: int, + buffer_size_mb: int, +) -> tuple[tuple[int, ...], list[tuple[slice, ...]]]: + if math.prod(shape) == 0: + return shape, [] + unit_shape = target_shards if target_shards is not None else target_chunks + copy_shape = compute_copy_shape( + shape, unit_shape, itemsize, max_buffer_bytes=buffer_size_mb * 1024 * 1024 + ) + slices_per_dim = [ + [slice(i, min(i + step, s)) for i in range(0, s, step)] + for s, step in zip(shape, copy_shape, strict=True) + ] + return copy_shape, list(itertools.product(*slices_per_dim)) + + +def require_zarr(tool_name: str) -> Any: + """Ensure the zarr package is installed.""" + try: + import zarr + + return zarr + except ImportError as exc: + raise ImportError(f"{tool_name} requires zarr; install it with: pip install 'blosc2[zarr]'") from exc + + +def copy_batches( + src: Any, + dst: Any, + batch_slices: list[tuple[slice, ...]], + total_bytes: int, + copy_direction: str, + quiet: bool = False, + verbose: bool = False, +) -> float: + """Copy data slices from ``src`` to ``dst`` (Blosc2 or Zarr arrays) with progress tracking.""" + total_batches = len(batch_slices) + if total_batches == 0: + return 0.0 + t_start = time.perf_counter() + + if HAVE_RICH and not quiet: + console = Console() + with Progress( + TextColumn("[bold blue]{task.description}"), + BarColumn(), + TaskProgressColumn(), + TransferSpeedColumn(), + TimeElapsedColumn(), + TimeRemainingColumn(), + console=console, + ) as progress: + task = progress.add_task(f"Converting {copy_direction}...", total=total_bytes) + for s in batch_slices: + data = src[s] + dst[s] = data + batch_size = data.nbytes + progress.update(task, advance=batch_size) + del data + else: + for idx, s in enumerate(batch_slices): + data = src[s] + dst[s] = data + del data + if verbose and not quiet: + pct = (idx + 1) / total_batches * 100 + print(f"Batch {idx + 1}/{total_batches} ({pct:.1f}%) written") + + return time.perf_counter() - t_start + + +def _fast_sample_verify(src: Any, dst: Any, slices_list: list[tuple[slice, ...]], quiet: bool) -> None: + corner_indices = [ + tuple(0 for _ in src.shape), + tuple(-1 for _ in src.shape), + tuple(s // 2 for s in src.shape), + ] + for idx in corner_indices: + v_src = src[idx] + v_dst = dst[idx] + if not np.array_equal(v_src, v_dst, equal_nan=True): + raise ValueError(f"Verification failed at index {idx}: source={v_src} != destination={v_dst}") + + if slices_list: + for s in (slices_list[0], slices_list[-1]): + src_chunk = src[s] + dst_chunk = dst[s] + if not np.array_equal(src_chunk, dst_chunk, equal_nan=True): + raise ValueError(f"Verification failed for slice {s}") + if not quiet: + print("✓ Fast sample verification passed!") + + +def _full_verify(src: Any, dst: Any, slices_list: list[tuple[slice, ...]], quiet: bool) -> None: + if not quiet: + print("Performing full verification across all slices...") + for idx, s in enumerate(slices_list): + src_data = src[s] + dst_data = dst[s] + if not np.array_equal(src_data, dst_data, equal_nan=True): + raise ValueError(f"Full verification failed at slice {s} (batch {idx})") + if not quiet: + print("✓ Full element verification passed!") + + +def verify_arrays( + src: Any, + dst: Any, + slices_list: list[tuple[slice, ...]], + full: bool = False, + quiet: bool = False, +) -> bool: + """Verify that data in the destination array matches the source array.""" + if src.shape != dst.shape: + raise ValueError(f"Shape mismatch: source {src.shape} vs destination {dst.shape}") + if src.dtype != dst.dtype: + raise ValueError(f"Dtype mismatch: source {src.dtype} vs destination {dst.dtype}") + + if math.prod(src.shape) == 0: + if not quiet: + print("✓ Empty array verification passed!") + return True + + if not full: + _fast_sample_verify(src, dst, slices_list, quiet) + else: + _full_verify(src, dst, slices_list, quiet) + return True + + +def print_summary(summary: dict[str, Any], dst_label: str) -> None: + """Print the conversion summary, using a rich table when available.""" + if HAVE_RICH: + console = Console() + table = Table(title="[bold green]Conversion Summary[/bold green]", box=None) + table.add_column("Property", style="bold cyan") + table.add_column("Value") + table.add_row("Source", f"{summary['src_path']} ({format_bytes(summary['src_size_bytes'])})") + table.add_row("Destination", f"{summary['dst_path']} ({format_bytes(summary['dst_size_bytes'])})") + table.add_row("Shape / Dtype", f"{summary['shape']} / {summary['dtype']}") + table.add_row("Chunks", str(summary["chunks"])) + if summary.get("shards") is not None: + table.add_row("Shards", str(summary["shards"])) + table.add_row("Uncompressed Size", format_bytes(summary["uncompressed_bytes"])) + table.add_row("Source Compression", f"{summary['src_cratio']:.2f}x") + table.add_row(f"{dst_label} Compression", f"{summary['dst_cratio']:.2f}x") + table.add_row( + "Conversion Speed", + f"{summary['throughput_mb_s']:.1f} MB/s ({summary['copy_duration_sec']:.2f} s)", + ) + table.add_row("Total Time", f"{summary['total_duration_sec']:.2f} s") + console.print() + console.print(table) + else: + print("\n--- Conversion Summary ---") + print(f"Source: {summary['src_path']} ({format_bytes(summary['src_size_bytes'])})") + print(f"Destination: {summary['dst_path']} ({format_bytes(summary['dst_size_bytes'])})") + print(f"Shape / Dtype: {summary['shape']} / {summary['dtype']}") + print(f"Chunks: {summary['chunks']}") + if summary.get("shards") is not None: + print(f"Shards: {summary['shards']}") + print(f"Uncompressed: {format_bytes(summary['uncompressed_bytes'])}") + print(f"Source CRatio: {summary['src_cratio']:.2f}x") + print(f"{dst_label + ' CRatio:':<18}{summary['dst_cratio']:.2f}x") + print( + f"Conversion Speed: {summary['throughput_mb_s']:.1f} MB/s ({summary['copy_duration_sec']:.2f} s)" + ) + print(f"Total Time: {summary['total_duration_sec']:.2f} s") diff --git a/src/blosc2/cli/zarr_to_b2nd.py b/src/blosc2/cli/zarr_to_b2nd.py new file mode 100644 index 000000000..0d0d6345e --- /dev/null +++ b/src/blosc2/cli/zarr_to_b2nd.py @@ -0,0 +1,399 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +"""Convert a Zarr array into a Blosc2 .b2nd file (NDArray).""" + +from __future__ import annotations + +import argparse +import contextlib +import math +import shutil +import sys +import time +from pathlib import Path +from typing import Any + +import numpy as np + +import blosc2 +from blosc2.cli.conversion_utils import ( + BLOSC2_CODECS, + BLOSC2_SHUFFLE_FILTERS, + DEFAULT_BUFFER_SIZE_MB, + compute_batches, + copy_batches, + get_path_size, + parse_shape, + print_summary, + require_zarr, + verify_arrays, +) + + +def zarr_to_b2nd( + input_path: str | Path, + output_path=None, + *, + chunks: tuple[int, ...] | None = None, + blocks: tuple[int, ...] | None = None, + codec: str = "auto", + clevel: int | None = None, + shuffle: str = "auto", + buffer_size_mb: int = DEFAULT_BUFFER_SIZE_MB, + overwrite: bool = False, + verify: bool = False, + full_verify: bool = False, + quiet: bool = False, + verbose: bool = False, +) -> dict[str, Any]: + zarr = require_zarr("zarr-to-b2nd") + src_path = Path(input_path).expanduser().resolve() + if not src_path.exists(): + raise FileNotFoundError(f"Input file not found: {src_path}") + + if output_path is None: + stem = src_path.name.removesuffix(".zip").removesuffix(".zarr") + dst_path = src_path.with_name(f"{stem}.b2nd") + else: + dst_path = Path(output_path).expanduser().resolve() + if src_path == dst_path or src_path.is_relative_to(dst_path) or dst_path.is_relative_to(src_path): + raise ValueError("Destination must not overlap the source path") + + if dst_path.exists() and not overwrite: + raise FileExistsError( + f"Destination path already exists: {dst_path}. Use -f/--force/--overwrite to replace it." + ) + + t_start = time.perf_counter() + z_arr, zip_store = open_zarr_input(zarr, src_path) + try: + shape = z_arr.shape + dtype = np.dtype(z_arr.dtype) + itemsize = dtype.itemsize + uncompressed_bytes = math.prod(shape) * itemsize + src_size = get_path_size(src_path) + + chunks, blocks = _determine_target_chunks_and_blocks(z_arr, chunks, blocks) + cparams = determine_blosc2_cparams(z_arr, codec, clevel, shuffle) + _, batch_slices = compute_batches(shape, chunks, None, itemsize, buffer_size_mb) + + # All validation is done, so only now remove the old output. + if dst_path.is_dir(): + shutil.rmtree(dst_path) + else: + dst_path.unlink(missing_ok=True) + + b2_arr = blosc2.empty( + shape, + dtype=dtype, + chunks=chunks, + blocks=blocks, + cparams=cparams, + urlpath=str(dst_path), + mode="w", + ) + copy_duration = copy_batches( + z_arr, b2_arr, batch_slices, uncompressed_bytes, "zarr -> b2nd", quiet=quiet, verbose=verbose + ) + copy_attrs(z_arr, b2_arr) + + if verify or full_verify: + verify_arrays(z_arr, b2_arr, batch_slices, full=full_verify, quiet=quiet) + finally: + if zip_store is not None: + zip_store.close() + + dst_size = get_path_size(dst_path) + total_duration = time.perf_counter() - t_start + throughput_mb_s = (uncompressed_bytes / (1024 * 1024)) / copy_duration if copy_duration > 0 else 0.0 + + return { + "src_path": src_path, + "dst_path": dst_path, + "shape": shape, + "dtype": dtype, + "chunks": chunks, + "blocks": blocks, + "cparams": cparams, + "src_size_bytes": src_size, + "dst_size_bytes": dst_size, + "uncompressed_bytes": uncompressed_bytes, + "copy_duration_sec": copy_duration, + "total_duration_sec": total_duration, + "throughput_mb_s": throughput_mb_s, + "src_cratio": uncompressed_bytes / src_size if src_size > 0 else 1.0, + "dst_cratio": uncompressed_bytes / dst_size if dst_size > 0 else 1.0, + } + + +def open_zarr_input(zarr: Any, path: Path) -> tuple[Any, Any]: + """Open a Zarr array for reading; return it and the ZipStore to close (or None).""" + zip_store = None + if path.name.endswith(".zip"): + zip_store = zarr.storage.ZipStore(str(path), mode="r") + store = zip_store if zip_store is not None else str(path) + + try: + z_arr = zarr.open_array(store=store, mode="r") + except Exception as exc: + is_group = type(exc).__name__ in {"ContainsGroupError", "NodeTypeValidationError"} + if type(exc).__name__ == "ArrayNotFoundError": + # Zarr v2 reports a missing array even when a group exists. + with contextlib.suppress(FileNotFoundError): + zarr.open_group(store=store, mode="r") + is_group = True + if zip_store is not None: + zip_store.close() + if is_group: + raise ValueError(f"{path} is a Zarr group; pass the path of an array") from exc + raise + return z_arr, zip_store + + +def _determine_target_chunks_and_blocks( + z, + chunks: tuple[int, ...] | None, + blocks: tuple[int, ...] | None, +) -> tuple[tuple[int, ...], tuple[int, ...]]: + """Choose b2nd chunks and blocks for a Zarr array. + + Zarr shards (or chunks, if unsharded) become b2nd chunks; Zarr inner chunks become blocks. + """ + shape = tuple(z.shape) + dtype = np.dtype(z.dtype) + + if z.shards is None: + target_chunks = chunks if chunks is not None else z.chunks + target_blocks = blocks + else: + target_chunks = chunks if chunks is not None else z.shards + target_blocks = blocks + if chunks is None and blocks is None: + target_blocks = z.chunks + + target_chunks = tuple(max(1, c) for c in target_chunks) + + # Fills in missing blocks, or checks that the given ones fit inside the chunks. + target_chunks, target_blocks = blosc2.compute_chunks_blocks( + shape, chunks=target_chunks, blocks=target_blocks, dtype=dtype + ) + + chunk_nbytes = math.prod(target_chunks) * dtype.itemsize + if chunk_nbytes > blosc2.MAX_BUFFERSIZE: + raise ValueError( + f"Chunk shape {tuple(target_chunks)} is {chunk_nbytes} bytes, larger than the " + f"Blosc2 limit of {blosc2.MAX_BUFFERSIZE} bytes; pass smaller --chunks" + ) + return tuple(target_chunks), tuple(target_blocks) + + +def determine_blosc2_cparams( + z_arr: Any, + codec: str = "auto", + clevel: int | None = None, + shuffle: str = "auto", +) -> blosc2.CParams: + """Translate the Zarr array's compressor into Blosc2 compression parameters.""" + compressors = z_arr.compressors + compressor = compressors[0] if compressors else None + kind = type(compressor).__name__ + itemsize = np.dtype(z_arr.dtype).itemsize + + if kind == "BloscCodec": + src_cname, src_clevel, src_shuffle = compressor.cname, compressor.clevel, compressor.shuffle + elif kind == "Blosc": + # numcodecs shuffle is an int; -1 (auto) means bitshuffle for 1-byte items. + src_cname, src_clevel = compressor.cname, compressor.clevel + src_shuffle = {0: "noshuffle", 1: "shuffle", 2: "bitshuffle"}.get( + compressor.shuffle, "bitshuffle" if itemsize == 1 else "shuffle" + ) + elif kind in {"ZstdCodec", "Zstd"}: + # Zstd level 0 means its default level (not "no compression"); levels go up to 22. + src_cname = "zstd" + src_clevel = min(compressor.level, 9) if compressor.level > 0 else 5 + src_shuffle = "shuffle" + elif kind in {"GzipCodec", "GZip"}: + # zlib is Blosc2's closest match to gzip. + src_cname, src_clevel, src_shuffle = "zlib", compressor.level, "shuffle" + else: + src_cname, src_clevel, src_shuffle = "zstd", 5, "shuffle" + + final_cname = src_cname if codec == "auto" else codec + final_clevel = src_clevel if clevel is None else clevel + final_shuffle = src_shuffle if shuffle == "auto" else shuffle + if not 0 <= final_clevel <= 9: + raise ValueError(f"Compression level must be between 0 and 9, got {final_clevel}") + + blosc2_codec = BLOSC2_CODECS.get(final_cname.lower()) + if blosc2_codec is None: + if codec != "auto": + raise ValueError(f"Unsupported codec '{codec}'. Supported: auto, {', '.join(BLOSC2_CODECS)}") + blosc2_codec = blosc2.Codec.ZSTD + blosc2_filter = BLOSC2_SHUFFLE_FILTERS.get(final_shuffle.lower()) + if blosc2_filter is None: + raise ValueError( + f"Unsupported shuffle '{final_shuffle}'. Supported: auto, {', '.join(BLOSC2_SHUFFLE_FILTERS)}" + ) + + return blosc2.CParams( + codec=blosc2_codec, clevel=final_clevel, filters=[blosc2_filter], typesize=itemsize + ) + + +def copy_attrs(z_arr: Any, b2_arr: blosc2.NDArray) -> None: + """Copy the Zarr array's user attributes into the Blosc2 array's vlmeta.""" + for key, value in z_arr.attrs.items(): + with contextlib.suppress(Exception): + b2_arr.schunk.vlmeta[key] = value + + +def build_parser(prog: str = "zarr-to-b2nd") -> argparse.ArgumentParser: + """Build the command-line argument parser.""" + return argparse.ArgumentParser( + prog=prog, + description="Convert a Zarr array (v2 or v3) into a Blosc2 .b2nd file (NDArray).", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=f"""\ +Examples: + {prog} data.zarr + {prog} data.zarr out.b2nd --chunks 100,1000 + {prog} data.zarr.zip -f --verify --full-verify + {prog} data.zarr --codec zstd --clevel 3 +""", + ) + + +def main(argv: list[str] | None = None) -> int: + """Main CLI entry point.""" + prog = Path(sys.argv[0]).name if sys.argv and sys.argv[0] else "zarr-to-b2nd" + parser = build_parser(prog=prog) + + parser.add_argument("input", help="Path to input Zarr array (directory or .zip)") + + parser.add_argument( + "output", + nargs="?", + default=None, + help="Path to output .b2nd file (default: .b2nd)", + ) + parser.add_argument( + "-o", + "--output-path", + dest="output_opt", + default=None, + help="Optional explicit output path (overrides positional output)", + ) + parser.add_argument( + "-c", + "--chunks", + type=str, + default=None, + help="Blosc2 chunk shape, e.g. '10,1000,1000' (default: the Zarr shards, or chunks if unsharded)", + ) + parser.add_argument( + "--blocks", + type=str, + default=None, + help="Blosc2 block shape (default: the Zarr inner chunks if sharded, otherwise computed)", + ) + parser.add_argument( + "--codec", + type=str, + default="auto", + choices=["auto", *BLOSC2_CODECS], + help="Compression codec for Blosc2 (default: 'auto', translated from the Zarr compressor)", + ) + parser.add_argument( + "--clevel", + type=int, + default=None, + help="Compression level (0-9, default: inherit from the Zarr compressor)", + ) + parser.add_argument( + "--shuffle", + type=str, + default="auto", + choices=["auto", "noshuffle", "shuffle", "bitshuffle"], + help="Shuffle filter (default: 'auto', inherit from the Zarr compressor)", + ) + parser.add_argument( + "-b", + "--buffer-size", + type=int, + default=DEFAULT_BUFFER_SIZE_MB, + help=f"Max memory buffer size in MB for batch copying (default: {DEFAULT_BUFFER_SIZE_MB})", + ) + parser.add_argument( + "-f", + "--force", + "--overwrite", + dest="overwrite", + action="store_true", + help="Overwrite destination store if it already exists", + ) + parser.add_argument( + "--verify", + action="store_true", + help="Verify array integrity (corner and sample equality check)", + ) + parser.add_argument( + "--full-verify", + action="store_true", + help="Perform thorough element-by-element verification of all chunks", + ) + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help="Show detailed conversion information", + ) + parser.add_argument( + "-q", + "--quiet", + action="store_true", + help="Suppress all progress and informational output", + ) + + args = parser.parse_args(argv) + + out_path = args.output_opt if args.output_opt is not None else args.output + parsed_chunks = parse_shape(args.chunks) if args.chunks else None + parsed_blocks = parse_shape(args.blocks) if args.blocks else None + + try: + summary = zarr_to_b2nd( + input_path=args.input, + output_path=out_path, + chunks=parsed_chunks, + blocks=parsed_blocks, + codec=args.codec, + clevel=args.clevel, + shuffle=args.shuffle, + buffer_size_mb=args.buffer_size, + overwrite=args.overwrite, + verify=args.verify, + full_verify=args.full_verify, + quiet=args.quiet, + verbose=args.verbose, + ) + except Exception as exc: + print(f"Error: {exc}", file=sys.stderr) + if args.verbose: + import traceback + + traceback.print_exc() + return 1 + + if not args.quiet: + print_summary(summary, dst_label="Blosc2") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_zarr_to_b2nd.py b/tests/test_zarr_to_b2nd.py new file mode 100644 index 000000000..e41c734ef --- /dev/null +++ b/tests/test_zarr_to_b2nd.py @@ -0,0 +1,294 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +from __future__ import annotations + +import numpy as np +import pytest + +import blosc2 +from blosc2.cli.b2nd_to_zarr import b2nd_to_zarr +from blosc2.cli.zarr_to_b2nd import main, zarr_to_b2nd + +zarr = pytest.importorskip("zarr") +numcodecs = pytest.importorskip("numcodecs") + + +def make_zarr(path, data, **kwargs): + z = zarr.create_array(store=str(path), shape=data.shape, dtype=data.dtype, **kwargs) + z[...] = data + return z + + +def test_zarr_to_b2nd_basic(tmp_path): + src = tmp_path / "data.zarr" + dst = tmp_path / "data.b2nd" + + shape = (100, 50) + data = np.arange(5000, dtype=np.float32).reshape(shape) + z = make_zarr(src, data, chunks=(20, 50)) + z.attrs["custom_key"] = "custom_value" + + summary = zarr_to_b2nd(src, verify=True) + assert summary["dst_path"] == dst + assert summary["shape"] == shape + assert summary["dtype"] == np.float32 + assert summary["chunks"] == (20, 50) + + b = blosc2.open(dst) + assert isinstance(b, blosc2.NDArray) + assert b.chunks == (20, 50) + np.testing.assert_array_equal(b[:], data) + assert b.schunk.vlmeta["custom_key"] == "custom_value" + + +def test_zarr_to_b2nd_v2_blosc_compressor_is_preserved(tmp_path): + src = tmp_path / "v2.zarr" + data = np.arange(1200, dtype=np.int32).reshape(30, 40) + compressor = numcodecs.Blosc(cname="lz4", clevel=7, shuffle=numcodecs.Blosc.BITSHUFFLE) + make_zarr(src, data, chunks=(10, 40), zarr_format=2, compressors=compressor) + + summary = zarr_to_b2nd(src, full_verify=True) + b = blosc2.open(summary["dst_path"]) + assert b.cparams.codec == blosc2.Codec.LZ4 + assert b.cparams.clevel == 7 + assert b.cparams.filters[0] == blosc2.Filter.BITSHUFFLE + np.testing.assert_array_equal(b[:], data) + + +@pytest.mark.parametrize( + ("kwargs", "codec", "clevel", "filt"), + [ + # Zarr's default is Zstd level 0, meaning "Zstd default", not "no compression". + ({}, blosc2.Codec.ZSTD, 5, blosc2.Filter.SHUFFLE), + ({"compressors": zarr.codecs.ZstdCodec(level=15)}, blosc2.Codec.ZSTD, 9, blosc2.Filter.SHUFFLE), + ({"compressors": zarr.codecs.GzipCodec(level=4)}, blosc2.Codec.ZLIB, 4, blosc2.Filter.SHUFFLE), + ({"compressors": None}, blosc2.Codec.ZSTD, 5, blosc2.Filter.SHUFFLE), + ( + {"compressors": zarr.codecs.BloscCodec(cname="zstd", clevel=3, shuffle="noshuffle")}, + blosc2.Codec.ZSTD, + 3, + blosc2.Filter.NOFILTER, + ), + ], + ids=["zstd-default", "zstd-high-level", "gzip", "uncompressed", "blosc-v3"], +) +def test_zarr_to_b2nd_compressor_translation(tmp_path, kwargs, codec, clevel, filt): + src = tmp_path / "data.zarr" + data = np.arange(1000, dtype=np.float64) + make_zarr(src, data, chunks=(100,), **kwargs) + + summary = zarr_to_b2nd(src) + cparams = summary["cparams"] + assert cparams.codec == codec + assert cparams.clevel == clevel + assert cparams.filters[0] == filt + np.testing.assert_array_equal(blosc2.open(summary["dst_path"])[:], data) + + +def test_zarr_to_b2nd_v2_autoshuffle_depends_on_itemsize(tmp_path): + compressor = numcodecs.Blosc(shuffle=numcodecs.Blosc.AUTOSHUFFLE) + for dtype, expected in [(np.uint8, blosc2.Filter.BITSHUFFLE), (np.float64, blosc2.Filter.SHUFFLE)]: + src = tmp_path / f"{np.dtype(dtype).name}.zarr" + make_zarr(src, np.arange(100, dtype=dtype), zarr_format=2, compressors=compressor) + assert zarr_to_b2nd(src)["cparams"].filters[0] == expected + + +def test_zarr_to_b2nd_sharded(tmp_path): + src = tmp_path / "sharded.zarr" + data = np.arange(1600, dtype=np.int64).reshape(40, 40) + make_zarr(src, data, chunks=(5, 40), shards=(20, 40)) + + summary = zarr_to_b2nd(src, verify=True) + assert summary["chunks"] == (20, 40) + assert summary["blocks"] == (5, 40) + + b = blosc2.open(summary["dst_path"]) + assert b.chunks == (20, 40) + assert b.blocks == (5, 40) + np.testing.assert_array_equal(b[:], data) + + +def test_zarr_to_b2nd_custom_options(tmp_path): + src = tmp_path / "cube.zarr" + dst = tmp_path / "custom.b2nd" + data = np.arange(30 * 40 * 50, dtype=np.int32).reshape(30, 40, 50) + make_zarr(src, data, chunks=(10, 10, 10)) + + summary = zarr_to_b2nd( + src, + dst, + chunks=(15, 20, 50), + blocks=(5, 10, 50), + codec="lz4", + clevel=3, + shuffle="bitshuffle", + full_verify=True, + ) + assert summary["dst_path"] == dst + + b = blosc2.open(dst) + assert b.chunks == (15, 20, 50) + assert b.blocks == (5, 10, 50) + assert b.cparams.codec == blosc2.Codec.LZ4 + assert b.cparams.clevel == 3 + assert b.cparams.filters[0] == blosc2.Filter.BITSHUFFLE + np.testing.assert_array_equal(b[:], data) + + +def test_roundtrip_b2nd_to_zarr_to_b2nd(tmp_path): + original = tmp_path / "original.b2nd" + zarr_path = tmp_path / "middle.zarr" + roundtrip = tmp_path / "roundtrip.b2nd" + + data = np.linspace(0, 1, 40 * 60, dtype=np.float64).reshape(40, 60) + a = blosc2.asarray(data, chunks=(20, 60), blocks=(5, 60), urlpath=original, mode="w") + a.schunk.vlmeta["units"] = "K" + + b2nd_to_zarr(original, zarr_path, sharded=True) + zarr_to_b2nd(zarr_path, roundtrip, full_verify=True) + + b = blosc2.open(roundtrip) + assert b.chunks == a.chunks + assert b.blocks == a.blocks + assert b.schunk.vlmeta["units"] == "K" + np.testing.assert_array_equal(b[:], data) + + +def test_zarr_to_b2nd_zip_input(tmp_path): + original = tmp_path / "arr.b2nd" + src = tmp_path / "arr.zarr.zip" + data = np.arange(1000, dtype=np.float64) + blosc2.asarray(data, chunks=(100,), urlpath=original, mode="w") + b2nd_to_zarr(original, src) + original.unlink() + + summary = zarr_to_b2nd(src, verify=True) + # data.zarr.zip should become data.b2nd, not data.zarr.b2nd + assert summary["dst_path"] == tmp_path / "arr.b2nd" + np.testing.assert_array_equal(blosc2.open(summary["dst_path"])[:], data) + + +def test_zarr_to_b2nd_rejects_group(tmp_path): + src = tmp_path / "group.zarr" + group = zarr.open_group(str(src), mode="w") + group.create_array("temperature", shape=(4,), dtype="i4") + + with pytest.raises(ValueError, match="Zarr group"): + zarr_to_b2nd(src) + assert not (tmp_path / "group.b2nd").exists() + + +def test_zarr_to_b2nd_overwrite_protection(tmp_path): + src = tmp_path / "data.zarr" + dst = tmp_path / "data.b2nd" + make_zarr(src, np.arange(10, dtype=np.int32)) + + zarr_to_b2nd(src, dst) + with pytest.raises(FileExistsError): + zarr_to_b2nd(src, dst, overwrite=False) + zarr_to_b2nd(src, dst, overwrite=True) + assert dst.exists() + + +def test_zarr_to_b2nd_cannot_write_inside_source(tmp_path): + src = tmp_path / "data.zarr" + data = np.arange(10, dtype=np.int32) + make_zarr(src, data) + + with pytest.raises(ValueError, match="source"): + zarr_to_b2nd(src, src / "inside.b2nd", overwrite=True) + np.testing.assert_array_equal(zarr.open_array(str(src), mode="r")[:], data) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"chunks": (10,)}, "same length"), + ({"blocks": (500, 20)}, "greater than chunks"), + ({"clevel": 12}, "between 0 and 9"), + ({"codec": "zstdd"}, "Unsupported codec"), + ({"shuffle": "bytes"}, "Unsupported shuffle"), + ], +) +def test_invalid_conversion_preserves_destination(tmp_path, kwargs, match): + src = tmp_path / "data.zarr" + dst = tmp_path / "data.b2nd" + make_zarr(src, np.arange(400, dtype=np.int32).reshape(20, 20), chunks=(10, 20)) + dst.write_bytes(b"preserve me") + + with pytest.raises(ValueError, match=match): + zarr_to_b2nd(src, dst, overwrite=True, **kwargs) + assert dst.read_bytes() == b"preserve me" + + +def test_zarr_to_b2nd_empty_array(tmp_path): + src = tmp_path / "empty.zarr" + zarr.create_array(store=str(src), shape=(0, 100), dtype="f4") + + summary = zarr_to_b2nd(src, verify=True) + assert summary["shape"] == (0, 100) + assert all(c > 0 for c in summary["chunks"]) + assert blosc2.open(summary["dst_path"]).shape == (0, 100) + + +def test_zarr_to_b2nd_with_nans(tmp_path): + src = tmp_path / "nans.zarr" + data = np.array([1.0, np.nan, 3.0, np.nan, 5.0], dtype=np.float32) + make_zarr(src, data, chunks=(2,)) + + summary = zarr_to_b2nd(src, full_verify=True) + np.testing.assert_array_equal(blosc2.open(summary["dst_path"])[:], data) + + +def test_zarr_to_b2nd_skips_unstorable_attrs(tmp_path): + src = tmp_path / "attrs.zarr" + z = make_zarr(src, np.arange(4, dtype=np.int32)) + # JSON allows integers of any size, but vlmeta only stores 64-bit integers. + z.attrs.update({"units": "K", "huge": 2**70, "nested": {"a": [1, 2]}}) + + summary = zarr_to_b2nd(src) + b = blosc2.open(summary["dst_path"]) # keep the array alive while reading its vlmeta + vlmeta = b.schunk.vlmeta + assert vlmeta["units"] == "K" + assert vlmeta["nested"] == {"a": [1, 2]} + assert "huge" not in vlmeta + + +def test_cli_main(tmp_path): + src = tmp_path / "cli.zarr" + dst = tmp_path / "cli.b2nd" + data = np.arange(50, dtype=np.int16) + make_zarr(src, data, chunks=(10,)) + + ret = main([str(src), str(dst), "--verify", "--chunks", "25", "--blocks", "5", "--quiet"]) + assert ret == 0 + + b = blosc2.open(dst) + assert b.chunks == (25,) + assert b.blocks == (5,) + np.testing.assert_array_equal(b[:], data) + + +def test_cli_main_plain_summary(tmp_path, capsys, monkeypatch): + src = tmp_path / "plain.zarr" + make_zarr(src, np.arange(50, dtype=np.int16)) + monkeypatch.setattr("blosc2.cli.conversion_utils.HAVE_RICH", False) + + assert main([str(src)]) == 0 + out = capsys.readouterr().out + assert "--- Conversion Summary ---" in out + assert "Blosc2 CRatio:" in out + + +def test_cli_main_errors(tmp_path, capsys): + assert main([str(tmp_path / "missing.zarr")]) == 1 + assert "Error:" in capsys.readouterr().err + + with pytest.raises(SystemExit) as excinfo: + main([str(tmp_path / "missing.zarr"), "--codec", "gzip"]) + assert excinfo.value.code == 2