Coverage for core/src/sensorkit/api/declarative.py: 87%

1# SPDX-License-Identifier: Apache-2.0

2from __future__ import annotations

4import asyncio

5import functools

6import inspect

7import warnings

8from collections.abc import Callable, Coroutine

9from enum import StrEnum, auto

10from typing import Any, final, get_type_hints, override

12from loguru import logger

14from sensorkit.api.bootstrap import connect

15from sensorkit.common.importutil import get_caller_module

16from sensorkit.core.client import SensorKit, ServiceContext

17from sensorkit.core.controller import TaskHandlerCallback

18from sensorkit.core.delegate import (

19 ControllerDelegate,

20 DeviceDelegate,

21 EntityDelegate,

22 ProgramDelegate,

23)

24from sensorkit.core.device import CommandHandlerCallback, DeviceCommand

25from sensorkit.core.entity import DeviceDetails

26from sensorkit.core.executor import TaskFactoryFunc

27from sensorkit.core.impl.controller import ControllerImpl

28from sensorkit.core.impl.device import DeviceImpl

29from sensorkit.core.impl.entity import EntityImpl

30from sensorkit.core.impl.program import ProgramImpl

31from sensorkit.core.task import Task

32from sensorkit.core.trait import Archetype, Trait

34type AnyEntityDecl = DeclaredEntity | DeclaredDevice | DeclaredController | DeclaredProgram

35type InitDeinitCallback = Callable[[], Coroutine[Any, Any, None] | None]

37AUTO_ENTITY_ATTR = "__sk_entity__"

38DECL_MARK_ATTR = "__sk_decl__"

39CALLBACK_MARK_ATTR = "__sk_callback__"

40TRAIT_ANNOTATION_ATTR = "__sk_traits__"

43def _mark_decl_type(cls: type, decl_type: type[DeclaredEntity]):

44 setattr(cls, DECL_MARK_ATTR, decl_type)

45 return cls

48def auto_create_decl(instance):

49 if decl_type := getattr(instance, DECL_MARK_ATTR, None):

50 decl = decl_type(None)

51 setattr(instance, AUTO_ENTITY_ATTR, decl)

53 # Pick up advisory trait annotations for devices.

54 if isinstance(decl, DeclaredDevice):

55 if traits := getattr(instance, TRAIT_ANNOTATION_ATTR, None):

56 decl._declared_traits = list(traits)

58 return instance

61def decl_for_instance(instance) -> AnyEntityDecl:

62 return getattr(instance, AUTO_ENTITY_ATTR, None)

65def entity_for_instance(instance):

66 return decl.impl if (decl := decl_for_instance(instance)) else None

69def _mark_callback(func: Callable, kind: CallbackKind):

70 """Mark a function as an unassociated callback."""

71 if inspect.ismethod(func):

72 # If we're working with a method, we need to unwrap it to get the underlying function. This

73 # is crucial because function objects are mutable (unlike methods) and are shared across

74 # subclass relationships.

75 func = func.__func__

77 setattr(func, CALLBACK_MARK_ATTR, kind)

80def is_callback(func: Callable):

81 """Return True if the given function is declared as a callback."""

82 return hasattr(func, CALLBACK_MARK_ATTR)

85def get_callback_kind(func: Callable):

86 """Return the type code of the callback function."""

87 return getattr(func, CALLBACK_MARK_ATTR)

90def introspect_param_type(func: Callable):

91 """Return the parameter type hint of the given single-parameter function."""

92 type_hints = get_type_hints(func)

93 type_hints.pop("return", None)

95 if len(type_hints) != 1:

96 raise DeclarationError("Callback must have exactly one typed parameter")

98 return next(iter(type_hints.values()))

100

101def introspect_decls(obj):

102 """Discover declared entities and callbacks within the given object."""

103 decls: list[DeclaredEntity] = []

104 callbacks: list[tuple[Callable, CallbackKind]] = []

105

106 for symbol in dir(obj):

107 if symbol.startswith("__") and symbol != AUTO_ENTITY_ATTR:

108 continue

109

110 # We use getattr_static here to avoid triggering property evaluations. This is quite slow,

111 # but with expected usage this performance should be acceptable.

112 attr = inspect.getattr_static(obj, symbol)

113

114 match attr:

115 case DeclaredEntity():

116 decls.append(attr)

117 case func if is_callback(func):

118 # Since getattr_static returns the function object, we need to do another getattr

119 # to get the method.

120 func = getattr(obj, symbol)

121 callbacks.append((func, get_callback_kind(func)))

122

123 return decls, callbacks

124

125

126class CallbackKind(StrEnum):

127 ENTITY_INIT = auto()

128 ENTITY_DEINIT = auto()

129 COMMAND_HANDLER = auto()

130 TASK_HANDLER = auto()

131 TASK_FACTORY = auto()

132 ENABLE = auto()

133 DISABLE = auto()

134

135

136class DeclaredEntity[T: EntityImpl = EntityImpl](EntityDelegate):

137 """Represents an eventual entity registration, collecting callbacks before the service starts."""

138

139 def __init__(self, name: str | None):

140 self.name = name

141 self._associated_callbacks: list[tuple[CallbackKind, Callable]] = []

142

143 # These are set when the service is registered.

144 self.client: SensorKit | None = None # TODO: Remove when the impl itself provides this.

145 self.service: ServiceContext | None = None

146 self.impl: T | None = None

147

148 @property

149 def delegate_target(self):

150 return self.impl

151

152 @property

153 def binding(self):

154 warnings.warn("Deprecated impl access from declarative object", stacklevel=2)

155 return self.impl

156

157 @final

158 def associate(self, func: Callable, kind: CallbackKind):

159 """Queue a callback to be registered with the implementation after it is created."""

160 self._associated_callbacks.append((kind, func))

161

162 @final

163 async def register(

164 self,

165 client: SensorKit,

166 service: ServiceContext,

167 *,

168 acquire_lease: bool = True,

169 ):

170 """Register the declared entity with the SensorKit backend.

171

172 Args:

173 client: the SensorKit client instance for interacting with the backend

174 service: the service context that manages this entity's lifecycle

175 acquire_lease: whether to acquire a lease on the entity during registration

176

177 Raises:

178 DeclarationError: if the entity does not have a name assigned

179 Exception: any exception raised by an initialization callback is propagated

180 """

181 if not self.name:

182 raise DeclarationError("Entity must have a name")

183

184 # Create our entity implementation and register all associated callbacks. Nothing hits the

185 # wire at this point.

186 self.impl = self.create_impl(service)

187

188 for kind, func in self._associated_callbacks:

189 self.register_callback(kind, func)

190

191 # Validate the declaration against the now-configured impl, before we touch the backend.

192 self.validate_declaration()

193

194 # Store references for use by the caller.

195 self.client = client

196 self.service = service

197

198 # Register and attach the impl. If lease acquisition fails, we raise here and the service

199 # will exit.

200 await service.register_impl(self.impl, acquire_lease=acquire_lease)

201

202 def validate_declaration(self):

203 """Validate the declaration against the configured implementation."""

204 pass

205

206 def create_impl(self, service: ServiceContext):

207 """Instantiate the implementation object bound to *service*."""

208 return EntityImpl.for_service_context(service, self.name)

209

210 def register_callback(self, kind: CallbackKind, func: Callable):

211 """Route a callback to the appropriate registration method on the implementation."""

212 match kind:

213 case CallbackKind.ENTITY_INIT:

214 self.impl.on_attach(func)

215 case CallbackKind.ENTITY_DEINIT:

216 self.impl.on_detach(func)

217

218

219class DeclaredDevice(DeclaredEntity[DeviceImpl], DeviceDelegate):

220 """Represents an eventual device registration."""

221

222 def __init__(self, name: str | None):

223 super().__init__(name)

224 self._declared_traits: list[Trait] = []

225

226 @override

227 def create_impl(self, svc: ServiceContext):

228 return DeviceImpl.for_service_context(svc, self.name)

229

230 @override

231 def validate_declaration(self):

232 # FIXME: For now we trust the device to publish the keywords required by its traits.

233 # This should be removed in favor of API that allows the device implementation

234 # to explicitly declare the keywords it publishes, which can then be used to

235 # validate whether its traits are satisfied.

236 for trait in self._declared_traits:

237 for kw_id in trait.effective_keyword_ids():

238 self.impl.declare_published_keyword(kw_id)

239

240 # Validate declared traits against the info the device will publish when it attaches.

241 details = self.impl.entity_info().details

242

243 if not isinstance(details, DeviceDetails):

244 raise RuntimeError("Device did not publish its details")

245

246 for trait in self._declared_traits:

247 if not trait.match(details):

248 missing = []

249

250 if commands := trait.effective_command_ids() - details.supported_commands:

251 missing.append("does not implement " + ", ".join(sorted(commands)))

252

253 if keywords := trait.effective_keyword_ids() - details.published_keywords:

254 missing.append("does not publish " + ", ".join(sorted(keywords)))

255

256 raise DeclarationError(

257 f"Device declares trait '{trait.name}' but {'; '.join(missing)}"

258 )

259

260 @override

261 def register_callback(self, kind: CallbackKind, func: Callable):

262 super().register_callback(kind, func)

263

264 match kind:

265 case CallbackKind.COMMAND_HANDLER:

266 command_type = introspect_param_type(func)

267

268 if not issubclass(command_type, DeviceCommand):

269 raise DeclarationError("Command handler parameter has incorrect type")

270

271 self.impl.command_handler(command_type)(func)

272

273

274class DeclaredController(DeclaredEntity[ControllerImpl], ControllerDelegate):

275 """Represents an eventual controller registration."""

276

277 @override

278 def create_impl(self, svc: ServiceContext):

279 return ControllerImpl.for_service_context(svc, self.name)

280

281 @override

282 def register_callback(self, kind: CallbackKind, func: Callable):

283 super().register_callback(kind, func)

284

285 match kind:

286 case CallbackKind.TASK_HANDLER:

287 task_type = introspect_param_type(func)

288

289 if not issubclass(task_type, Task):

290 raise DeclarationError("Task handler parameter has incorrect type")

291

292 self.impl.task_handler(task_type)(func)

293

294

295class DeclaredProgram(DeclaredEntity[ProgramImpl], ProgramDelegate):

296 """Represents an eventual program registration."""

297

298 @override

299 def create_impl(self, svc: ServiceContext):

300 return ProgramImpl.for_service_context(svc, self.name)

301

302 @override

303 def register_callback(self, kind: CallbackKind, func: Callable):

304 super().register_callback(kind, func)

305

306 match kind:

307 case CallbackKind.ENABLE:

308 self.impl.on_enable(func)

309 case CallbackKind.DISABLE:

310 self.impl.on_disable(func)

311 case CallbackKind.TASK_FACTORY:

312 self.impl.task_factory(func)

313

314

315class Service:

316 """Runs a single SensorKit service instance given a set of declared entity objects."""

317

318 def __init__(self, name: str, version: str):

319 self.name = name

320 self.version = version

321 self.declarations: set[DeclaredEntity] = set()

322 self.context: ServiceContext | None = None

323 self.client: SensorKit | None = None

324 self._delegate_entity: DeclaredEntity | None = None

325 self._register_lock = asyncio.Lock()

326 self._started = False

327 loop = asyncio.get_running_loop()

328 self.running = loop.create_future()

329 self.shutdown = loop.create_future()

330

331 def add(self, declared: DeclaredEntity, name: str | None = None):

332 """Add an entity declaration to the service.

333

334 >>> my_entity = declare_entity(name="my_entity")

335 >>> service.add(my_entity)

336

337 If `name` is given and the declaration already includes a static name assignment, an error

338 is raised.

339

340 If no `name` is given and the declaration has no static name assigned, the entity will

341 "share" the name of the containing service. This can only apply to the first such nameless

342 declaration. Adding a second declaration without a name will raise an error.

343 """

344 if name:

345 if declared.name:

346 raise DeclarationError("Cannot reassign entity name")

347

348 declared.name = name

349 elif not declared.name:

350 if self._delegate_entity:

351 raise DeclarationError("Cannot determine entity name")

352

353 # Set the name of this entity to the name of the service. This allows single-entity

354 # services to avoid having to name both the service and the entity.

355 self._delegate_entity = declared

356 declared.name = self.name

357

358 self.declarations.add(declared)

359

360 def include(self, obj: Any, *, name: str | None = None):

361 """Add all entity declarations contained or marked in the given object.

362

363 Typically, the target object will be the instance of a class that has been decorated by,

364 e.g., `declare_entity`. In this case, the marks applied by that decorator are found and

365 a single entity declaration is created automatically.

366

367 >>> @declare_entity

368 >>> class MyEntity:

369 >>> ...

370 >>> service.include(MyEntity(), name="my_entity")

371

372 If the input object is a type, the type will be instantiated assuming a no-parameter

373 constructor. The instance object will then be searched for marks and entity declarations

374 as above.

375

376 >>> service.include(MyEntity, name="my_entity")

377

378 If the input object is a module, marks do not apply. It is simply searched for entity

379 declarations.

380

381 >>> import myorg.devices

382 >>> service.include(myorg.devices)

383

384 In all cases, the `name` argument is applied if and only if a single entity declaration is

385 resolved. If more than one is resolved (e.g., in the module case), an error is raised to

386 indicate the name assignment ambiguity.

387 """

388 if isinstance(obj, type):

389 cls = obj

390

391 try:

392 obj = cls()

393 except Exception as e:

394 raise DeclarationError(f"Failed to instantiate class {cls.__name__}") from e

395

396 typename = type(obj).__name__

397

398 # If we have an instance of a class marked for automatic entity creation, make it so.

399 auto_create_decl(obj)

400

401 # Get the declarations contained in the input object.

402 decls, callbacks = introspect_decls(obj)

403

404 if not decls:

405 raise DeclarationError(f"No entity declarations found in {typename}")

406

407 if len(decls) == 1:

408 # Exactly one declaration in this namespace. This will typically be the case for

409 # classes that implement an entity.

410 decl = decls[0]

411

412 # Automatically associate all floating callbacks and then add the declaration.

413 for func, kind in callbacks:

414 decl.associate(func, kind)

415

416 self.add(decl, name)

417 else:

418 # Multiple declarations. This will generally be a module include. In this case, we

419 # treat the existence of any floating callbacks as an error condition rather than make

420 # assumptions about which declaration is intended to be associated with which callback.

421 if callbacks:

422 raise DeclarationError(

423 f"Floating callbacks alongside multiple declarations in {typename}"

424 )

425

426 # Similarly, we can't be sure about name assignment either. The module include use case

427 # demands hardcoded naming.

428 if name is not None:

429 raise DeclarationError(

430 f"Name assignment is ambiguous with multiple declarations in {typename}"

431 )

432

433 # Add all declarations.

434 for decl in decls:

435 self.add(decl)

436

437 return decls

438

439 def include_module(self, **kwargs):

440 """Add all entity declarations found in the calling module.

441

442 See the `include` method.

443 """

444 # Include decls from the module of the user code call site. We must check for None here

445 # because in certain contexts there may be no calling module.

446 if mod := get_caller_module(depth=1):

447 self.include(mod, **kwargs)

448

449 async def register(self):

450 """Idempotent method to connect to the backend and register this service."""

451 async with self._register_lock:

452 if self.client is None:

453 self.client = await connect()

454

455 if self.context is None:

456 self.context = await self.client.register_service(self.name, self.version)

457

458 async def start(self):

459 """Start the service."""

460 if self._started:

461 raise RuntimeError("Service was already started")

462

463 self._started = True

464

465 try:

466 # Register as a service.

467 await self.register()

468

469 # Register all declared entities.

470 await asyncio.gather(

471 *(

472 decl.register(

473 self.client, self.context, acquire_lease=decl is not self._delegate_entity

474 )

475 for decl in self.declarations

476 )

477 )

478

479 self.running.set_result(True)

480 except BaseException as e:

481 logger.debug(f"service error propagated to declarative API: {type(e).__name__}: {e}")

482 self.running.set_exception(e)

483 self.shutdown.set_exception(e)

484

485 if self.context is not None:

486 await self.context.shutdown()

487

488 raise

489

490 async def wait_for_shutdown():

491 try:

492 await self.context.join()

493 self.shutdown.set_result(True)

494 except asyncio.CancelledError:

495 self.shutdown.cancel()

496 raise

497 except BaseException as e:

498 if not self.shutdown.done():

499 self.shutdown.set_exception(e)

500

501 if not isinstance(e, Exception):

502 # Make sure to re-raise BaseException.

503 raise

504

505 self._shutdown_waiter = asyncio.create_task(wait_for_shutdown())

506

507 async def run(self):

508 """Run the service."""

509 await self.start()

510

511 try:

512 await self.shutdown

513 finally:

514 await self.stop()

515

516 async def stop(self):

517 """Stop the service."""

518 if self.context is not None:

519 await self.context.shutdown()

520

521 await self.shutdown

522

523

524class DeclarationError(Exception):

525 """Raised when a declaration or its usage is invalid."""

526

527

528def _decorator(decl: DeclaredEntity | None, kind: CallbackKind, func: Callable):

529 if decl:

530 # Associate the callback with the declaration.

531 decl.associate(func, kind)

532 else:

533 # Mark as floating, for later automatic association.

534 _mark_callback(func, kind)

535

536 return func

537

538

539def on_attach(arg: DeclaredEntity | InitDeinitCallback):

540 """Register a callback to be executed during entity initialization."""

541 match arg:

542 case DeclaredEntity():

543 return functools.partial(_decorator, arg, CallbackKind.ENTITY_INIT)

544 case _:

545 return _decorator(None, CallbackKind.ENTITY_INIT, arg)

546

547

548def on_detach(arg: DeclaredEntity | InitDeinitCallback):

549 """Register a callback to be executed during entity deinitialization."""

550 match arg:

551 case DeclaredEntity():

552 return functools.partial(_decorator, arg, CallbackKind.ENTITY_DEINIT)

553 case _:

554 return _decorator(None, CallbackKind.ENTITY_DEINIT, arg)

555

556

557def command_handler(arg: DeclaredEntity | CommandHandlerCallback):

558 """Declare a command handler."""

559 match arg:

560 case DeclaredEntity():

561 return functools.partial(_decorator, arg, CallbackKind.COMMAND_HANDLER)

562 case _:

563 return _decorator(None, CallbackKind.COMMAND_HANDLER, arg)

564

565

566def task_handler(arg: DeclaredEntity | TaskHandlerCallback):

567 """Declare a task handler for a controller."""

568 match arg:

569 case DeclaredEntity():

570 return functools.partial(_decorator, arg, CallbackKind.TASK_HANDLER)

571 case _:

572 return _decorator(None, CallbackKind.TASK_HANDLER, arg)

573

574

575def task_factory(arg: DeclaredEntity | TaskFactoryFunc):

576 """Declare the task factory function for a program."""

577 match arg:

578 case DeclaredEntity():

579 return functools.partial(_decorator, arg, CallbackKind.TASK_FACTORY)

580 case _:

581 return _decorator(None, CallbackKind.TASK_FACTORY, arg)

582

583

584def on_enable(arg: DeclaredEntity | Callable):

585 """Register a callback to be called when a program or device is enabled."""

586 match arg:

587 case DeclaredEntity():

588 return functools.partial(_decorator, arg, CallbackKind.ENABLE)

589 case _:

590 return _decorator(None, CallbackKind.ENABLE, arg)

591

592

593def on_disable(arg: DeclaredEntity | Callable):

594 """Register a callback to be called when a program or device is disabled."""

595 match arg:

596 case DeclaredEntity():

597 return functools.partial(_decorator, arg, CallbackKind.DISABLE)

598 case _:

599 return _decorator(None, CallbackKind.DISABLE, arg)

600

601

602def declare_entity(cls: type | None = None, *, name: str | None = None):

603 """Declare a generic entity to be implemented by a local service."""

604 if cls:

605 return _mark_decl_type(cls, DeclaredEntity)

606 else:

607 return DeclaredEntity(name)

608

609

610def declare_device(

611 cls: type | None = None,

612 *,

613 name: str | None = None,

614 type: Archetype | None = None,

615 traits: list[Trait] | None = None,

616):

617 """Declare a Device to be implemented by a local service.

618

619 Can be used as a bare class decorator (`@sk.declare_device`), a keyword

620 decorator (`@sk.declare_device(type=..., traits=[...])`), or an explicit declaration

621 (`my_device = sk.declare_device(name="foo", type=..., traits=[...])`)

622 """

623 all_traits: list[Trait] | None = None

624 if type is not None or traits:

625 all_traits = ([type] if type is not None else []) + list(traits or [])

626

627 if cls:

628 # Bare decorator: @sk.declare_device

629 if all_traits:

630 setattr(cls, TRAIT_ANNOTATION_ATTR, all_traits)

631 return _mark_decl_type(cls, DeclaredDevice)

632 elif name is not None:

633 # Explicit declaration: sk.declare_device(name="foo")

634 decl = DeclaredDevice(name)

635 if all_traits:

636 decl._declared_traits = all_traits

637 return decl

638 elif all_traits is not None:

639 # Keyword decorator: @sk.declare_device(type=..., traits=[...])

640 def decorator(decorated_cls):

641 setattr(decorated_cls, TRAIT_ANNOTATION_ATTR, all_traits)

642 return _mark_decl_type(decorated_cls, DeclaredDevice)

643 return decorator

644 else:

645 return DeclaredDevice(None)

646

647

648def declare_controller(cls: type | None = None, *, name: str | None = None):

649 """Declare a Controller to be implemented by a local service."""

650 if cls:

651 return _mark_decl_type(cls, DeclaredController)

652 else:

653 return DeclaredController(name)

654

655

656def declare_program(cls: type | None = None, *, name: str | None = None):

657 """Declare a Program to be implemented by a local service."""

658 if cls:

659 return _mark_decl_type(cls, DeclaredProgram)

660 else:

661 return DeclaredProgram(name)

Coverage for core / src / sensorkit / api / declarative.py: 87%

350 statements