|
6 | 6 | from functools import cached_property |
7 | 7 | from itertools import chain |
8 | 8 | from pathlib import Path |
9 | | -from typing import Self |
| 9 | +from typing import Any, Self |
10 | 10 |
|
11 | 11 | from bluesky_stomp.messaging import MessageContext, StompClient |
12 | 12 | from bluesky_stomp.models import Broker |
|
38 | 38 | ) |
39 | 39 | from blueapi.utils import deprecated |
40 | 40 | from blueapi.worker import WorkerEvent, WorkerState |
41 | | -from blueapi.worker.event import ProgressEvent, TaskStatus |
| 41 | +from blueapi.worker.event import ProgressEvent, TaskError, TaskResult, TaskStatus |
42 | 42 | from blueapi.worker.task_worker import TrackableTask |
43 | 43 |
|
44 | 44 | from .event_bus import AnyEvent, EventBusClient, OnAnyEvent |
@@ -141,13 +141,17 @@ def __init__(self, name, model: PlanModel, client: "BlueapiClient"): |
141 | 141 | self._client = client |
142 | 142 | self.__doc__ = model.description |
143 | 143 |
|
144 | | - def __call__(self, *args, **kwargs): |
| 144 | + def __call__(self, *args, **kwargs) -> Any: |
145 | 145 | req = TaskRequest( |
146 | 146 | name=self.name, |
147 | 147 | params=self._build_args(*args, **kwargs), |
148 | 148 | instrument_session=self._client.instrument_session, |
149 | 149 | ) |
150 | | - self._client.run_task(req) |
| 150 | + match self._client.run_task(req): |
| 151 | + case TaskStatus(result=TaskResult(result=res)): |
| 152 | + return res |
| 153 | + case TaskStatus(result=TaskError(type=typ, message=msg)): |
| 154 | + raise PlanFailedError(typ, msg) |
151 | 155 |
|
152 | 156 | @property |
153 | 157 | def help_text(self) -> str: |
@@ -744,3 +748,9 @@ def login(self, token_path: Path | None = None): |
744 | 748 | auth.start_device_flow() |
745 | 749 | else: |
746 | 750 | print("Server is not configured to use authentication!") |
| 751 | + |
| 752 | + |
| 753 | +class PlanFailedError(Exception): |
| 754 | + def __init__(self, typ: str, message: str): |
| 755 | + super().__init__(message) |
| 756 | + self._type = typ |
0 commit comments