Coverage for core / src / sensorkit / core / trait.py: 100%

1# SPDX-License-Identifier: Apache-2.0 

2"""Device trait infrastructure. 

3 

4Traits are high-level interfaces that describe device capabilities by specifying 

5which commands a device must implement. They use structural typing: a device matches 

6a trait if it implements all required commands. 

7 

8Two variants exist: 

9 - Trait: a device can satisfy many; specifies only required commands. 

10 - Archetype (subclass of Trait): a device should match at most one; additionally 

11 supports required and optional sub-traits, and optional commands. 

12""" 

13 

14from __future__ import annotations 

15 

16import itertools 

17from collections.abc import Iterable 

18from dataclasses import dataclass 

19from typing import TYPE_CHECKING 

20 

21if TYPE_CHECKING: 

22 from sensorkit.core.device import DeviceCommand 

23 from sensorkit.core.entity import DeviceDetails 

24 

25# Global registries, auto-populated by declare_trait() and declare_archetype(). 

26_trait_registry: set[Trait] = set() 

27_archetype_registry: set[Archetype] = set() 

28 

29 

30@dataclass(frozen=True, eq=False) 

31class Trait: 

32 """A named set of required commands and/or keywords that a device may structurally satisfy.""" 

33 

34 name: str 

35 required_commands: tuple[type[DeviceCommand], ...] = () 

36 required_keywords: tuple[str, ...] = () 

37 

38 def effective_command_ids(self) -> frozenset[str]: 

39 """Return all command IDs required by this trait.""" 

40 return frozenset(cmd.model_tag() for cmd in self.required_commands) 

41 

42 def effective_keyword_ids(self) -> frozenset[str]: 

43 """Return all keyword IDs required by this trait.""" 

44 return frozenset(self.required_keywords) 

45 

46 def match(self, details: DeviceDetails) -> bool: 

47 """Return True if the given device details satisfy this trait's commands and keywords.""" 

48 if not (self.effective_command_ids() <= details.supported_commands): 

49 return False 

50 if not (self.effective_keyword_ids() <= details.published_keywords): 

51 return False 

52 return True 

53 

54 def __hash__(self): 

55 return hash(self.name) 

56 

57 def __eq__(self, other): 

58 return isinstance(other, Trait) and self.name == other.name 

59 

60 def __repr__(self): 

61 return f"Trait({self.name!r})" 

62 

63 

64@dataclass(frozen=True, eq=False) 

65class Archetype(Trait): 

66 """A trait archetype: a device should match at most one. 

67 

68 Extends Trait with support for required and optional sub-traits, and optional commands. 

69 Optional traits and commands do not affect matching but describe extended capabilities. 

70 """ 

71 

72 required_traits: tuple[Trait, ...] = () 

73 optional_commands: tuple[type[DeviceCommand], ...] = () 

74 

75 def effective_command_ids(self) -> frozenset[str]: 

76 """Return all command IDs required by this archetype and its required sub-traits.""" 

77 return frozenset( 

78 itertools.chain( 

79 (cmd.model_tag() for cmd in self.required_commands), 

80 *(sub_trait.effective_command_ids() for sub_trait in self.required_traits), 

81 ) 

82 ) 

83 

84 def effective_keyword_ids(self) -> frozenset[str]: 

85 """Return all keyword IDs required by this archetype and its required sub-traits.""" 

86 return frozenset( 

87 itertools.chain( 

88 self.required_keywords, 

89 *(sub_trait.effective_keyword_ids() for sub_trait in self.required_traits), 

90 ) 

91 ) 

92 

93 def __repr__(self): 

94 return f"Archetype({self.name!r})" 

95 

96 

97def declare_trait( 

98 name: str, 

99 *, 

100 required_commands: tuple[type[DeviceCommand], ...] = (), 

101 required_keywords: tuple[str, ...] = (), 

102) -> Trait: 

103 """Declare a device trait and add it to the global registry.""" 

104 trait = Trait( 

105 name=name, 

106 required_commands=required_commands, 

107 required_keywords=required_keywords, 

108 ) 

109 _trait_registry.add(trait) 

110 return trait 

111 

112 

113def declare_archetype( 

114 name: str, 

115 *, 

116 required_commands: tuple[type[DeviceCommand], ...] = (), 

117 required_keywords: tuple[str, ...] = (), 

118 required_traits: tuple[Trait, ...] = (), 

119 optional_commands: tuple[type[DeviceCommand], ...] = (), 

120) -> Archetype: 

121 """Declare an archetype and add it to the global registries. 

122 

123 Archetypes extend traits with required and optional sub-traits, and optional commands. 

124 Sub-traits must themselves be plain Traits, not Archetypes. 

125 """ 

126 for t in required_traits: 

127 if isinstance(t, Archetype): 

128 raise ValueError( 

129 f"Sub-trait '{t.name}' is an archetype and cannot be used as a sub-trait" 

130 ) 

131 

132 archetype = Archetype( 

133 name=name, 

134 required_commands=required_commands, 

135 required_keywords=required_keywords, 

136 required_traits=required_traits, 

137 optional_commands=optional_commands, 

138 ) 

139 

140 _trait_registry.add(archetype) 

141 _archetype_registry.add(archetype) 

142 

143 return archetype 

144 

145 

146def get_registered_traits() -> frozenset[Trait]: 

147 """Return all registered traits (including archetypes).""" 

148 return frozenset(_trait_registry) 

149 

150 

151def get_registered_archetypes() -> frozenset[Archetype]: 

152 """Return all registered archetypes.""" 

153 return frozenset(_archetype_registry) 

154 

155 

156def match_traits( 

157 details: DeviceDetails, 

158 traits: Iterable[Trait] | None = None, 

159 *, 

160 exclude_archetypes: bool = False, 

161) -> list[Trait]: 

162 """Return all traits from the given iterable that the device satisfies.""" 

163 if traits is None: 

164 traits = _trait_registry 

165 

166 output = [] 

167 

168 for trait in traits: 

169 if exclude_archetypes and isinstance(trait, Archetype): 

170 continue 

171 

172 if trait.match(details): 

173 output.append(trait) 

174 

175 return output 

176 

177 

178def match_archetype(details: DeviceDetails) -> Archetype | None: 

179 """Return the first matching archetype for the given device details, or None.""" 

180 for archetype in _archetype_registry: 

181 if archetype.match(details): 

182 return archetype 

183 

184 return None