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

@@ -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")

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)
 

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,17 +18,18 @@
 
 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,
  EnvironmentResponse,
  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()