diff --git a/server/cli/annotate.py b/server/cli/annotate.py index abaac16a..c5fb90ff 100644 --- a/server/cli/annotate.py +++ b/server/cli/annotate.py @@ -4,7 +4,7 @@ import os.path import shlex import shutil import subprocess -import sys +from os.path import isfile from subprocess import STDOUT, PIPE from tempfile import NamedTemporaryFile @@ -27,23 +27,42 @@ def annotate_args(func): @sort_options @click.command( - short_help="Annotate H5AD file columns. Run `cellxgene annotation --help` for more information.", - options_metavar="", + options_metavar="" ) -@click.option( - "-i", - "--input-h5ad-file", +@click.argument( + "input_h5ad_file", + type=click.Path(exists=True, dir_okay=False, readable=True), + nargs=1, + metavar="", required=True, - type=str, - help="The input H5AD file containing the missing annotations.", ) @click.option( "-m", "--model-url", + # Making this a required "option", rather than an "argument", since we support automatic model selection in the + # future, in which case the user would not need to specify this option at all and we can make it optional at + # that time. required=True, help="The URL of the model used to prediction annotated labels. May be a local filesystem directory " "or S3 path (s3://)", ) +@click.option( + "-o", + "--output-h5ad-file", + default="", + help="The output H5AD file that will contain the generated annotation values. If this option is not provided, " + "the input file will be overwritten to include the new annotations; in this case you must specify " + "--overwrite.", + metavar="", +) +@click.option( + "--overwrite", + default=False, + is_flag=True, + help="Allow overwriting of the specified H5AD output file, if it exists. For safety, you must specify this " + "flag if the specified output file already exists or if the --output-h5ad-file option is not provided.", + show_default=True, +) @click.option( "-l", "--counts-layer", @@ -53,8 +72,8 @@ def annotate_args(func): @click.option( "-g", "--gene-column-name", - help="The name of the `var` column that contains gene identifiers. The values in this column will be used to match " - "genes between the query and reference datasets. If not specified, the gene identifiers are expected to exist " + help="The name of the `var` column that contains gene names. The values in this column will be used to match " + "genes between the query and reference datasets. If not specified, the gene names are expected to exist " "in `var.index`.", ) # TODO: Useful if we want to support discoverability of models @@ -91,19 +110,6 @@ def annotate_args(func): "will store the predicted annotation values and confidence scores. This can be used to allow multiple " "annotation predictions to be run on a single AnnData object.", ) -@click.option( - "-u", - "--update-h5ad-file", - is_flag=True, - help="Flag indicating whether to update the input h5ad file with annotation values. This option is mutually " - "exclusive with --output-h5ad-file.", -) -@click.option( - "-o", - "--output-h5ad-file", - help="The output H5AD file that will contain the generated annotation values. This option is mutually " - "exclusive with --update-h5ad-file.", -) @click.option("--use-model-cache/--no-use-model-cache", default=True) @click.option( "--use-gpu/--no-use-gpu", @@ -141,6 +147,9 @@ def annotate_args(func): ) @click.help_option("--help", "-h", help="Show this message and exit.") def annotate(**cli_args): + """ + Add predicted annotations to an H5AD file. Run `cellxgene annotate --help` for more information. + """ _validate_options(cli_args) print(f"Reading query dataset {cli_args['input_h5ad_file']}...") @@ -149,7 +158,11 @@ def annotate(**cli_args): filter(None, [cli_args.get("annotation_prefix"), cli_args.get("annotation_type"), cli_args.get("run_name")]) ) - output_h5ad_file = cli_args["input_h5ad_file"] if cli_args["update_h5ad_file"] else cli_args["output_h5ad_file"] + output_h5ad_file = ( + cli_args["input_h5ad_file"] + if cli_args["overwrite"] and not cli_args["output_h5ad_file"] + else cli_args["output_h5ad_file"] + ) model_url = cli_args.get("model_url") local_model_path = _retrieve_model(cli_args.get("model_cache_dir"), model_url, cli_args.get("use_model_cache")) @@ -196,7 +209,7 @@ def annotate(**cli_args): p.wait() if p.returncode == 0: - print(f"Wrote annotations to {cli_args.get('output_h5ad_file')}") + print(f"Wrote annotations to {output_h5ad_file}") else: print("Annotation failed!") else: @@ -218,13 +231,11 @@ def _retrieve_model(model_cache_dir, model_url, use_cache=True): def _validate_options(cli_args): - # TODO(atolopko): Use cloup library for this logic - if cli_args["update_h5ad_file"] and cli_args["output_h5ad_file"]: - click.echo("--update_h5ad_file and --output_h5ad_file are mutually exclusive") - sys.exit(1) - if not (cli_args["update_h5ad_file"] or cli_args["output_h5ad_file"]): - click.echo("--update_h5ad_file or --output_h5ad_file must be specified") - sys.exit(1) + output = cli_args["output_h5ad_file"] + overwrite = cli_args["overwrite"] + + if isfile(output) and not overwrite: + raise click.UsageError(f"Cannot overwrite existing file {output}, try using the flag --overwrite") if __name__ == "__main__": diff --git a/test/unit/cli/test_annotate.py b/test/unit/cli/test_annotate.py index c8fe1b4b..bd75edf4 100644 --- a/test/unit/cli/test_annotate.py +++ b/test/unit/cli/test_annotate.py @@ -1,7 +1,7 @@ import os import shutil import unittest -from tempfile import mkstemp, TemporaryDirectory +from tempfile import mkstemp, TemporaryDirectory, NamedTemporaryFile import mlflow from click.testing import CliRunner @@ -12,10 +12,8 @@ from test.unit.cli.fixtures.mlflow_model_fixture import FakeModel def write_model(model) -> str: with TemporaryDirectory() as mlflow_model_dir: - fixtures_path = os.path.join(os.path.dirname(__file__), 'fixtures') - mlflow.pyfunc.save_model(mlflow_model_dir, - loader_module='fixtures', - code_path=[fixtures_path]) + fixtures_path = os.path.join(os.path.dirname(__file__), "fixtures") + mlflow.pyfunc.save_model(mlflow_model_dir, loader_module="fixtures", code_path=[fixtures_path]) return shutil.make_archive(mkstemp()[1], "zip", mlflow_model_dir) @@ -46,7 +44,6 @@ class TestCliAnnotate(unittest.TestCase): result = CliRunner().invoke( annotate, [ - "--input-h5ad-file", query_dataset_file_path, "--model-url", model_file_path, @@ -54,7 +51,8 @@ class TestCliAnnotate(unittest.TestCase): f"{query_dataset_file_path}.output", # avoid having mflow create conda env or virtualenv when in test env; # this avoids making pip remote requests and is also faster - "--mlflow-env-manager", "local" + "--mlflow-env-manager", + "local", ], ) @@ -74,31 +72,52 @@ class TestCliAnnotate(unittest.TestCase): result.stdout, "inputs passed correctly", ) - - def test__annotate__verifies_mutually_exclusive_options(self): - required_options = ["--input-h5ad-file", "some.h5ad", "--model-url", "some_url"] - result = CliRunner().invoke( - annotate, - required_options + [], - ) - - self.assertNotEqual(0, result.exit_code, "aborts with non-success code") self.assertIn( - "--update_h5ad_file or --output_h5ad_file must be specified", + f"Wrote annotations to {query_dataset_file_path}.output", result.stdout, - "error message displayed", + "success message is correct", ) - result = CliRunner().invoke( - annotate, required_options + ["--output-h5ad-file", "some_arg", "--update-h5ad-file"] - ) + def test__annotate__requires_overwrite_option_when_output_file_exists(self): - self.assertNotEqual(0, result.exit_code, "aborts with non-success code") - self.assertIn( - "--update_h5ad_file and --output_h5ad_file are mutually exclusive", - result.stdout, - "error message displayed", - ) + with NamedTemporaryFile() as input_h5ad, NamedTemporaryFile() as existing_file: + required_options = [input_h5ad.name, "--output-h5ad-file", existing_file.name, "--model-url", "some_url"] + result = CliRunner().invoke( + annotate, + required_options + [], + ) + + self.assertNotEqual(0, result.exit_code, "aborts with non-success code") + self.assertIn( + "try using the flag --overwrite", + result.stdout, + "error message displayed", + ) + + def test__annotate__overwrite_option_allows_overwrite_of_existing_output_file(self): + model_file_path = write_model(FakeModel()) + + with NamedTemporaryFile() as existing_file: + required_options = [ + existing_file.name, + "--output-h5ad-file", + existing_file.name, + "--overwrite", + "--model-url", + model_file_path, + ] + result = CliRunner().invoke( + annotate, + required_options + [], + ) + + print(result.stdout) + self.assertNotEqual(1, result.exit_code, "aborts with non-success code") + self.assertIn( + f"Wrote annotations to {existing_file.name}", + result.stdout, + "success message is correct on output file overwrite", + ) # TODO: