diff --git a/.gitignore b/.gitignore index f67d2adb..bfdc38f2 100644 --- a/.gitignore +++ b/.gitignore @@ -25,7 +25,8 @@ share/python-wheels/ .installed.cfg *.egg MANIFEST - +pip-* +tmp* # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. @@ -150,4 +151,5 @@ cython_debug/ mdio1/* */mdio1/* pytest-of-* -tmp/ +tmp +debugging/* diff --git a/noxfile.py b/noxfile.py index ac271f2f..291167a5 100644 --- a/noxfile.py +++ b/noxfile.py @@ -132,6 +132,7 @@ def precommit(session: Session) -> None: "pre-commit", "pre-commit-hooks", "pyupgrade", + "pytest-dependency", ) session.run("pre-commit", *args) if args and args[0] == "install": diff --git a/src/mdio/__main__.py b/src/mdio/__main__.py index 10e3ae3e..5234ff73 100644 --- a/src/mdio/__main__.py +++ b/src/mdio/__main__.py @@ -10,6 +10,8 @@ import click_params import mdio +from mdio.commands.utility import copy +from mdio.commands.utility import info plugin_folder = os.path.join(os.path.dirname(__file__), "commands") @@ -40,11 +42,18 @@ class MyCLI(click.MultiCommand): http://lybniz2.sourceforge.net/safeeval.html """ + _command_mapping = { + "copy": copy, + "info": info, + } + + protected_files = ["__init__.py", "utility.py"] + def list_commands(self, ctx: click.Context) -> list[str]: """List commands available under `commands` module.""" - rv = [] + rv = list(self._command_mapping.keys()) for filename in os.listdir(plugin_folder): - if filename.endswith(".py") and filename != "__init__.py": + if filename.endswith(".py") and filename not in self.protected_files: rv.append(filename[:-3]) rv.sort() @@ -61,6 +70,8 @@ def get_command(self, ctx: click.Context, name: str) -> dict[Callable]: } local_ns = {} + if name in self._command_mapping.keys(): + return self._command_mapping[name] fn = os.path.join(plugin_folder, name + ".py") with open(fn) as f: code = compile(f.read(), fn, "exec") diff --git a/src/mdio/commands/utility.py b/src/mdio/commands/utility.py new file mode 100644 index 00000000..e3bb7628 --- /dev/null +++ b/src/mdio/commands/utility.py @@ -0,0 +1,177 @@ +"""Dataset CLI Plugin.""" + +try: + from typing import Optional + + import click + import click_params + + import mdio +except SystemError: + pass + + +@click.command(name="copy") +@click.option( + "-i", + "--input-mdio-path", + required=True, + help="Input mdio path.", + type=click.Path(exists=True), +) +@click.option( + "-o", + "--output-mdio-path", + required=True, + help="Output path or URL to write the mdio dataset.", + type=click.STRING, +) +@click.option( + "-access", + "--access-pattern", + required=False, + default="012", + help="Access pattern of the file", + type=click.STRING, + show_default=True, +) +@click.option( + "-exc", + "--excludes", + required=False, + default="", + help="""Data to exclude during copy. i.e. chunked_012. The raw data won’t be + copied, but it will create an empty array to be filled. If left blank, it will + copy everything.""", + type=click.STRING, +) +@click.option( + "-inc", + "--includes", + required=False, + default="", + help="""Data to include during copy. i.e. trace_headers. If this is not + specified, and certain data is excluded, it will not copy headers. To + preserve headers, specify trace_headers. If left blank, it will copy + everything except specified in excludes parameter.""", + type=click.STRING, +) +@click.option( + "-storage", + "--storage-options", + required=False, + help="Custom storage options for cloud backends", + type=click_params.JSON, +) +@click.option( + "-overwrite", + "--overwrite", + required=False, + default=False, + help="Flag to overwrite if mdio file if it exists", + type=click.BOOL, + show_default=True, +) +def copy( + input_mdio_path: str, + output_mdio_path: str, + access_pattern: str = "012", + includes: str = "", + excludes: str = "", + storage_options: Optional[dict] = None, + overwrite: bool = False, +): + """Copy a MDIO dataset to anpther MDIO dataset. + + Can also copy with empty data to be filled later. See `excludes` + and `includes` parameters. + + More documentation about `excludes` and `includes` can be found + in Zarr's documentation in `zarr.convenience.copy_store`. + """ + reader = mdio.MDIOReader( + input_mdio_path, access_pattern=access_pattern, return_metadata=True + ) + mdio.copy_mdio( + source=reader, + dest_path_or_buffer=output_mdio_path, + excludes=excludes, + includes=includes, + storage_options=storage_options, + overwrite=overwrite, + ) + + +@click.command(name="info") +@click.option( + "-i", + "--input-mdio-file", + required=True, + help="Input path of the mdio file", + type=click.STRING, +) +@click.option( + "-access", + "--access-pattern", + required=False, + default="012", + help="Access pattern of the file", + type=click.STRING, + show_default=True, +) +@click.option( + "-format", + "--output-format", + required=False, + default="plain", + help="""Output format, plain is human readable. JSON will output in json + format for easier passing. """, + type=click.Choice(["plain", "json"]), + show_default=True, + show_choices=True, +) +def info( + input_mdio_file, + output_format, + access_pattern, +): + """Provide information on a MDIO dataset. + + By default this returns human readable information about the grid and stats for + the dataset. If output-format is set to json then a json is returned to + facilitate parsing. + """ + reader = mdio.MDIOReader( + input_mdio_file, access_pattern=access_pattern, return_metadata=True + ) + mdio_dict = {} + mdio_dict["grid"] = {} + for axis in reader.grid.dim_names: + dim = reader.grid.select_dim(axis) + min = dim.coords[0] + max = dim.coords[-1] + size = dim.coords.shape[0] + axis_dict = {"name": axis, "min": min, "max": max, "size": size} + mdio_dict["grid"][axis] = axis_dict + + if output_format == "plain": + click.echo("{:<10} {:<10} {:<10} {:<10}".format("NAME", "MIN", "MAX", "SIZE")) + click.echo("=" * 40) + + for _, axis_dict in mdio_dict["grid"].items(): + click.echo( + "{:<10} {:<10} {:<10} {:<10}".format( + axis_dict["name"], + axis_dict["min"], + axis_dict["max"], + axis_dict["size"], + ) + ) + + click.echo("\n\n{:<10} {:<10}".format("STAT", "VALUE")) + click.echo("=" * 20) + for name, stat in reader.stats.items(): + click.echo(f"{name:<10} {stat:<10}") + if output_format == "json": + mdio_dict["stats"] = reader.stats + click.echo(mdio_dict) diff --git a/tests/conftest.py b/tests/conftest.py index 9015bd95..19e9e851 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -32,6 +32,13 @@ def zarr_tmp(tmp_path_factory): return tmp_file +@pytest.fixture(scope="session") +def zarr_tmp2(tmp_path_factory): + """Make a temp file for the output MDIO.""" + tmp_file = tmp_path_factory.mktemp(r"mdio2") + return tmp_file + + @pytest.fixture(scope="session") def segy_export_ibm_tmp(tmp_path_factory): """Make a temp file for the round-trip IBM SEG-Y.""" diff --git a/tests/test_main.py b/tests/test_main.py index ac91dfe5..87d2787d 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -15,6 +15,7 @@ def runner() -> CliRunner: return CliRunner() +@pytest.mark.dependency() def test_main_succeeds(runner: CliRunner, segy_input: str, zarr_tmp: Path) -> None: """It exits with a status code of zero.""" cli_args = ["segy", "import"] @@ -27,6 +28,27 @@ def test_main_succeeds(runner: CliRunner, segy_input: str, zarr_tmp: Path) -> No assert result.exit_code == 0 +@pytest.mark.dependency(depends=["test_main_succeeds"]) +def test_main_info_succeeds(runner: CliRunner, zarr_tmp: Path) -> None: + """It exits with a status code of zero.""" + cli_args = ["info"] + cli_args.extend(["-i", str(zarr_tmp)]) + + result = runner.invoke(__main__.main, args=cli_args) + assert result.exit_code == 0 + + +@pytest.mark.dependency(depends=["test_main_succeeds"]) +def test_main_copy_succeeds(runner: CliRunner, zarr_tmp: Path, zarr_tmp2: Path) -> None: + """It exits with a status code of zero.""" + cli_args = ["copy"] + cli_args.extend(["-i", str(zarr_tmp)]) + cli_args.extend(["-o", str(zarr_tmp2)]) + + result = runner.invoke(__main__.main, args=cli_args) + assert result.exit_code == 0 + + def test_cli_version(runner: CliRunner) -> None: """Check if version prints without error.""" cli_args = ["--version"]