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()