Coverage for core / src / sensorkit / core / device.py: 92%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import uuid 

5from abc import abstractmethod 

6from collections.abc import Coroutine, Iterable 

7from typing import TYPE_CHECKING, Any, Callable, ClassVar 

8 

9from pydantic import BaseModel 

10 

11from sensorkit.backend.event import Event 

12from sensorkit.backend.request import Call, ExtendedResponse, Request 

13from sensorkit.common.model import ModelRegistry, RegistryBaseModel 

14from sensorkit.core.entity import ( 

15 DeviceDetails, 

16 EntityClient, 

17 EntityInfo, 

18 EntityInterface, 

19 EntityRef, 

20) 

21from sensorkit.core.state import EventSourcedState 

22from sensorkit.core.trait import Trait, match_archetype, match_traits 

23 

24if TYPE_CHECKING: 

25 from sensorkit.core.client import SensorKit 

26 

27type CommandHandlerCallback = Callable[[DeviceCommand], Coroutine[Any, Any, BaseModel | None]] 

28 

29 

30class DeviceCommand(RegistryBaseModel): 

31 """Base Device Command model.""" 

32 command_id: str = None 

33 

34 registry: ClassVar[ModelRegistry[DeviceCommand]] = ModelRegistry(discriminator="command_id") 

35 

36 @classmethod 

37 def model_registry(cls): 

38 return cls.registry 

39 

40 

41class Abort(DeviceCommand): 

42 """Built-in abort command.""" 

43 

44 

45class CommandStarted(Event): 

46 """Event indicating a Command has been accepted and execution has begun.""" 

47 command_id: str 

48 call_id: uuid.UUID 

49 

50 

51class CommandDone(Event): 

52 """Event indicating a Command has finished execution.""" 

53 command_id: str 

54 call_id: uuid.UUID 

55 success: bool 

56 

57 

58class DeviceEnableState(Event): 

59 """Event indicating the Device enable state has changed.""" 

60 enabled: bool 

61 

62 

63class DeviceState(EventSourcedState): 

64 """Device state.""" 

65 enable_state: DeviceEnableState 

66 

67 

68class DeviceEnableStateRequest(BaseModel): 

69 """Request that a Device enable or disable its command handlers.""" 

70 enable: bool 

71 

72 

73set_enable_state_request = Request.define( 

74 "set_enable_state", 

75 payload=DeviceEnableStateRequest, 

76) 

77 

78 

79class CommandRequestMessage(BaseModel): 

80 """Message model representing a Command execution request.""" 

81 command: DeviceCommand 

82 

83 

84class CommandResult(BaseModel): 

85 """Message model representing a response to a Command execution request.""" 

86 data: Any 

87 

88 

89run_command_request = Request.define( 

90 "command", 

91 payload=CommandRequestMessage, 

92 result=CommandResult, 

93) 

94 

95 

96class DeviceClient(EntityClient): 

97 """Object that exposes client-side functionality of a Device.""" 

98 

99 def __init__(self, *args, **kwargs): 

100 super().__init__(*args, **kwargs) 

101 self._details: DeviceDetails | None = None 

102 

103 async def enable(self): 

104 """Enable device command handlers.""" 

105 return await self.call( 

106 set_enable_state_request, 

107 DeviceEnableStateRequest(enable=True) 

108 ) 

109 

110 async def disable(self): 

111 """Disable device command handlers.""" 

112 return await self.call( 

113 set_enable_state_request, 

114 DeviceEnableStateRequest(enable=False) 

115 ) 

116 

117 def command(self, command: DeviceCommand) -> Call[ExtendedResponse, CommandResult]: 

118 """Send a command to the device and return a Call tracking the result.""" 

119 return self.call( 

120 run_command_request, 

121 CommandRequestMessage(command=command), 

122 ) 

123 

124 async def get_details(self) -> DeviceDetails: 

125 """Fetch and cache the device's supported command details. 

126 

127 The result is cached because a device's supported commands do not change at 

128 runtime. 

129 """ 

130 if self._details is None: 

131 info = await self.kv_get_model(EntityInfo) 

132 if not isinstance(info.details, DeviceDetails): 

133 raise ValueError(f"Entity is not a device: {self.entity}") 

134 self._details = info.details 

135 return self._details 

136 

137 async def has_trait(self, trait: Trait) -> bool: 

138 """Return True if this device satisfies the given trait.""" 

139 details = await self.get_details() 

140 return trait.match(details) 

141 

142 async def get_traits(self, traits: Iterable[Trait]) -> list[Trait]: 

143 """Return all matching traits from the given iterable.""" 

144 details = await self.get_details() 

145 return match_traits(details, traits) 

146 

147 async def get_archetype(self) -> Trait | None: 

148 """Return the matching archetype for this device, if any.""" 

149 details = await self.get_details() 

150 return match_archetype(details) 

151 

152 

153class DeviceRef(EntityRef[DeviceClient]): 

154 """A serializable reference to a device client.""" 

155 

156 def _get_client(self, kit: SensorKit) -> DeviceClient: 

157 return kit.device(self.name) 

158 

159 

160class DeviceInterface(EntityInterface): 

161 """Interface describing an implementation of a device.""" 

162 

163 @abstractmethod 

164 def on_enable(self, func: Callable[[], None]): 

165 """Register a callback to invoke when the device is enabled.""" 

166 ... 

167 

168 @abstractmethod 

169 def on_disable(self, func: Callable[[], None]): 

170 """Register a callback to invoke when the device is disabled.""" 

171 ... 

172 

173 @abstractmethod 

174 def command_handler( 

175 self, 

176 command_type: type[DeviceCommand], 

177 ) -> Callable[..., CommandHandlerCallback]: 

178 """Register a handler for the given command type and return a decorator.""" 

179 ...