From 07c273f55d419664e11aa4eb2400e3ad3b5a5c4e Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Mon, 9 Sep 2024 17:08:19 +0100 Subject: [PATCH] Extend formatting options to listen subcommand (#539) Listen streams a mix of WorkerEvent, ProgressEvent and DataEvent instances for anything that the server is doing. This extends the type handling in the format to each of these types to display the relevant level of information for each. Informational messages are changed to be printed to stderr so that the main stdout output can be piped to an external application (eg jq for json mode). --- src/blueapi/cli/cli.py | 12 ++- src/blueapi/cli/format.py | 35 +++++++- tests/unit_tests/test_cli.py | 161 +++++++++++++++++++++++------------ 3 files changed, 149 insertions(+), 59 deletions(-) diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index 6455a9814..9a86b9a92 100644 --- a/src/blueapi/cli/cli.py +++ b/src/blueapi/cli/cli.py @@ -1,5 +1,6 @@ import json import logging +import sys from functools import wraps from pathlib import Path from pprint import pprint @@ -158,20 +159,23 @@ def listen_to_events(obj: dict) -> None: else: raise RuntimeError("Message bus needs to be configured") + fmt = obj["fmt"] + def on_event( event: WorkerEvent | ProgressEvent | DataEvent, context: MessageContext, ) -> None: - converted = json.dumps(event.model_dump(), indent=2) - print(converted) + fmt.display(event) print( "Subscribing to all bluesky events from " - f"{config.stomp.host}:{config.stomp.port}" + f"{config.stomp.host}:{config.stomp.port}", + file=sys.stderr, ) with event_bus_client: event_bus_client.subscribe_to_all_events(on_event) - input("Press enter to exit") + print("Press enter to exit", file=sys.stderr) + input() @controller.command(name="run") diff --git a/src/blueapi/cli/format.py b/src/blueapi/cli/format.py index 3baaafa99..8d8ac744a 100644 --- a/src/blueapi/cli/format.py +++ b/src/blueapi/cli/format.py @@ -10,13 +10,24 @@ from pydantic import BaseModel +from blueapi.core.bluesky_types import DataEvent from blueapi.service.model import DeviceResponse, PlanResponse +from blueapi.worker.event import ProgressEvent, WorkerEvent FALLBACK = pprint +NL = "\n" Stream = TextIO | None +def fmt_dict(t: dict[str, Any] | Any, ind: int = 1) -> str: + """Format a (possibly nested) dict into a human readable tree""" + if not isinstance(t, dict): + return f" {t}" + pre = " " * (ind * 4) + return NL + NL.join(f"{pre}{k}:{fmt_dict(v, ind+1)}" for k, v in t.items() if v) + + class OutputFormat(str, enum.Enum): JSON = "json" FULL = "full" @@ -49,6 +60,17 @@ def display_full(obj: Any, stream: Stream): print(dev.name) for proto in dev.protocols: print(" " + proto) + case DataEvent(name=name, doc=doc): + print(f"{name.title()}:{fmt_dict(doc)}") + case WorkerEvent(state=st, task_status=task): + print( + f"WorkerEvent: {st.name}{fmt_dict(task.model_dump() if task else {})}" + ) + case ProgressEvent(): + print(f"Progress:{fmt_dict(obj.model_dump())}") + case BaseModel(): + print(obj.__class__.__name__, end="") + print(fmt_dict(obj.model_dump())) case other: FALLBACK(other, stream=stream) @@ -61,7 +83,7 @@ def display_json(obj: Any, stream: Stream): case DeviceResponse(devices=devices): print(json.dumps([d.model_dump() for d in devices], indent=2)) case BaseModel(): - print(json.dumps(obj.model_dump(), indent=2)) + print(json.dumps(obj.model_dump())) case _: print(json.dumps(obj)) @@ -83,6 +105,17 @@ def display_compact(obj: Any, stream: Stream): for dev in devices: print(dev.name) print(indent(textwrap.fill(", ".join(dev.protocols), 80), " ")) + case DataEvent(name=name): + print(f"Data Event: {name}") + case WorkerEvent(state=state): + print(f"Worker Event: {state.name}") + case ProgressEvent(statuses=stats): + prog = ( + max(100 * (s.percentage or 0) for s in stats.values()) + if stats + else "???" + ) + print(f"Progress: {prog}%") case other: FALLBACK(other, stream=stream) diff --git a/tests/unit_tests/test_cli.py b/tests/unit_tests/test_cli.py index 7f9e4cabc..afc0cdece 100644 --- a/tests/unit_tests/test_cli.py +++ b/tests/unit_tests/test_cli.py @@ -4,6 +4,7 @@ from io import StringIO from pathlib import Path from textwrap import dedent +from typing import Any from unittest.mock import Mock, patch import pytest @@ -17,10 +18,10 @@ from blueapi import __version__ from blueapi.cli.cli import main -from blueapi.cli.format import OutputFormat +from blueapi.cli.format import OutputFormat, fmt_dict from blueapi.client.rest import BlueskyRemoteControlError from blueapi.config import ScratchConfig, ScratchRepository -from blueapi.core.bluesky_types import Plan +from blueapi.core.bluesky_types import DataEvent, Plan from blueapi.service.model import ( DeviceModel, DeviceResponse, @@ -28,6 +29,7 @@ PlanModel, PlanResponse, ) +from blueapi.worker.event import ProgressEvent, TaskStatus, WorkerEvent, WorkerState @pytest.fixture @@ -165,10 +167,10 @@ def test_valid_stomp_config_for_listener( ], input="\n", ) - assert ( - result.output - == "Subscribing to all bluesky events from localhost:61613\nPress enter to exit" - ) + assert result.output == dedent("""\ + Subscribing to all bluesky events from localhost:61613 + Press enter to exit + """) assert result.exit_code == 0 @@ -347,13 +349,8 @@ def test_device_output_formatting(): HasName """) - output = StringIO() - OutputFormat.COMPACT.display(devices, out=output) - assert output.getvalue() == compact + _assert_matching_formatting(OutputFormat.COMPACT, devices, compact) - # json outputs valid json - output = StringIO() - OutputFormat.JSON.display(devices, out=output) json_out = dedent("""\ [ { @@ -364,16 +361,14 @@ def test_device_output_formatting(): } ] """) - assert output.getvalue() == json_out - _ = json.loads(output.getvalue()) + _assert_matching_formatting(OutputFormat.JSON, devices, json_out) + _ = json.loads(json_out) - output = StringIO() - OutputFormat.FULL.display(devices, out=output) full = dedent("""\ my-device HasName """) - assert output.getvalue() == full + _assert_matching_formatting(OutputFormat.FULL, devices, full) class ExtendedModel(BaseModel): @@ -405,13 +400,8 @@ def test_plan_output_formatting(): metadata=object """) - output = StringIO() - OutputFormat.COMPACT.display(plans, out=output) - assert output.getvalue() == compact + _assert_matching_formatting(OutputFormat.COMPACT, plans, compact) - # json outputs valid json - output = StringIO() - OutputFormat.JSON.display(plans, out=output) json_out = dedent("""\ [ { @@ -456,11 +446,9 @@ def test_plan_output_formatting(): } ] """) - assert output.getvalue() == json_out - _ = json.loads(output.getvalue()) + _assert_matching_formatting(OutputFormat.JSON, plans, json_out) + _ = json.loads(json_out) - output = StringIO() - OutputFormat.FULL.display(plans, out=output) full = dedent("""\ my-plan Summary of description @@ -504,45 +492,104 @@ def test_plan_output_formatting(): "type": "object" } """) - assert output.getvalue() == full + _assert_matching_formatting(OutputFormat.FULL, plans, full) + + +def test_event_formatting(): + data = DataEvent( + name="start", doc={"foo": "bar", "fizz": {"buzz": (1, 2, 3), "hello": "world"}} + ) + worker = WorkerEvent( + state=WorkerState.RUNNING, + task_status=TaskStatus(task_id="count", task_complete=False, task_failed=False), + errors=[], + warnings=[], + ) + progress = ProgressEvent(task_id="start", statuses={}) + + _assert_matching_formatting( + OutputFormat.JSON, + data, + ( + """{"name": "start", "doc": """ + """{"foo": "bar", "fizz": {"buzz": [1, 2, 3], "hello": "world"}}}\n""" + ), + ) + _assert_matching_formatting(OutputFormat.COMPACT, data, "Data Event: start\n") + _assert_matching_formatting( + OutputFormat.FULL, + data, + dedent("""\ + Start: + foo: bar + fizz: + buzz: (1, 2, 3) + hello: world + """), + ) + + _assert_matching_formatting( + OutputFormat.JSON, + worker, + ( + """{"state": "RUNNING", "task_status": """ + """{"task_id": "count", "task_complete": false, "task_failed": false}, """ + """"errors": [], "warnings": []}\n""" + ), + ) + _assert_matching_formatting(OutputFormat.COMPACT, worker, "Worker Event: RUNNING\n") + _assert_matching_formatting( + OutputFormat.FULL, + worker, + "WorkerEvent: RUNNING\n task_id: count\n", + ) + + _assert_matching_formatting( + OutputFormat.JSON, progress, """{"task_id": "start", "statuses": {}}\n""" + ) + _assert_matching_formatting(OutputFormat.COMPACT, progress, "Progress: ???%\n") + _assert_matching_formatting( + OutputFormat.FULL, progress, "Progress:\n task_id: start\n" + ) def test_unknown_object_formatting(): demo = {"foo": 42, "bar": ["hello", "World"]} - output = StringIO() - OutputFormat.JSON.display(demo, output) exp = """{"foo": 42, "bar": ["hello", "World"]}\n""" - assert exp == output.getvalue() + _assert_matching_formatting(OutputFormat.JSON, demo, exp) - output = StringIO() - OutputFormat.COMPACT.display(demo, output) exp = """{'bar': ['hello', 'World'], 'foo': 42}\n""" - assert exp == output.getvalue() + _assert_matching_formatting(OutputFormat.COMPACT, demo, exp) - output = StringIO() - OutputFormat.FULL.display(demo, output) - assert exp == output.getvalue() + _assert_matching_formatting(OutputFormat.FULL, demo, exp) + + +def test_dict_formatting(): + demo = {"name": "foo", "keys": [1, 2, 3], "metadata": {"fizz": "buzz"}} + exp = """\nname: foo\nkeys: [1, 2, 3]\nmetadata:\n fizz: buzz""" + assert fmt_dict(demo, 0) == exp + + demo = "not a dict" + assert fmt_dict(demo, 0) == " not a dict" def test_generic_base_model_formatting(): - output = StringIO() - obj = ExtendedModel(name="foo", keys=[1, 2, 3], metadata={"fizz": "buzz"}) - exp = dedent("""\ - { - "name": "foo", - "keys": [ - 1, - 2, - 3 - ], - "metadata": { - "fizz": "buzz" - } - } - """) - OutputFormat.JSON.display(obj, output) - assert exp == output.getvalue() + model = ExtendedModel(name="demo", keys=[1, 2, 3], metadata={"fizz": "buzz"}) + exp = '{"name": "demo", "keys": [1, 2, 3], "metadata": {"fizz": "buzz"}}\n' + _assert_matching_formatting(OutputFormat.JSON, model, exp) + + _assert_matching_formatting( + OutputFormat.FULL, + model, + dedent("""\ + ExtendedModel + name: demo + keys: [1, 2, 3] + metadata: + fizz: buzz + """), + ) @patch("blueapi.cli.cli.setup_scratch") @@ -564,3 +611,9 @@ def test_init_scratch_calls_setup_scratch(mock_setup_scratch: Mock, runner: CliR ) assert result.exit_code == 0 mock_setup_scratch.assert_called_once_with(expected_config) + + +def _assert_matching_formatting(fmt: OutputFormat, obj: Any, expected: str): + output = StringIO() + fmt.display(obj, output) + assert expected == output.getvalue()