diff --git a/src/hipercow/cli.py b/src/hipercow/cli.py index 4744cbb..6cf5993 100644 --- a/src/hipercow/cli.py +++ b/src/hipercow/cli.py @@ -1,7 +1,10 @@ +from functools import reduce +from operator import ior + import click from hipercow import root -from hipercow.task import task_list, task_status +from hipercow.task import TaskStatus, task_list, task_status from hipercow.task_create import task_create from hipercow.task_eval import task_eval @@ -30,9 +33,11 @@ def cli_task_status(task_id: str): @task.command("list") -def cli_task_list(): +@click.option("--with-status", type=str, multiple=True) +def cli_task_list(with_status=None): r = root.open_root() - for task_id in task_list(r): + with_status = _process_with_status(with_status) + for task_id in task_list(r, with_status=with_status): click.echo(task_id) @@ -50,3 +55,9 @@ def cli_task_create(cmd: tuple[str]): def cli_task_eval(task_id: str): r = root.open_root() task_eval(r, task_id) + + +def _process_with_status(with_status: list[str]): + if not with_status: + return None + return reduce(ior, [TaskStatus[i.upper()] for i in with_status]) diff --git a/tests/test_cli.py b/tests/test_cli.py index b511d9c..26a1c9c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3,6 +3,7 @@ from click.testing import CliRunner from hipercow import cli, root, task +from hipercow.task import TaskStatus def test_can_init_repository(tmp_path): @@ -49,6 +50,15 @@ def test_can_run_task(tmp_path): assert task.task_status(r, task_id) == task.TaskStatus.SUCCESS +def test_can_process_with_status_args(): + assert cli._process_with_status([]) is None + assert cli._process_with_status(["success"]) == TaskStatus.SUCCESS + assert ( + cli._process_with_status(["success", "running"]) + == TaskStatus.RUNNING | TaskStatus.SUCCESS + ) + + def test_can_list_tasks(tmp_path): runner = CliRunner() with runner.isolated_filesystem(temp_dir=tmp_path): @@ -60,3 +70,11 @@ def test_can_list_tasks(tmp_path): res = runner.invoke(cli.cli_task_list, []) assert res.exit_code == 0 assert res.output.strip() == task_id + + res = runner.invoke(cli.cli_task_list, ["--with-status", "created"]) + assert res.exit_code == 0 + assert res.output.strip() == task_id + + res = runner.invoke(cli.cli_task_list, ["--with-status", "running"]) + assert res.exit_code == 0 + assert res.output.strip() == ""