Coverage for src/loman/nodekey.py: 100%

1"""Node key implementation for computation graph navigation.""" 

2 

3import json 

4import re 

5from collections.abc import Hashable, Iterable 

6from dataclasses import dataclass 

7from typing import ClassVar, Optional, Union 

8 

9from loman.util import as_iterable 

10 

11Name = Union[str, "NodeKey", Hashable] 

12Names = list[Name] 

13 

14 

15class PathNotFoundError(Exception): 

16 """Exception raised when a node path cannot be found.""" 

17 

18 

19def quote_part(part: Hashable) -> str: 

20 """Quote a node key part for safe representation in paths.""" 

21 if isinstance(part, str): 

22 if "/" in part: 

23 return json.dumps(part) 

24 else: 

25 return part 

26 return str(part) 

27 

28 

29@dataclass(frozen=True, repr=False) 

30class NodeKey: 

31 """Immutable key for identifying nodes in the computation graph hierarchy.""" 

32 

33 parts: tuple[Hashable, ...] 

34 

35 def __str__(self) -> str: 

36 """Return string representation using path notation.""" 

37 return "/".join([quote_part(part) for part in self.parts]) 

38 

39 @property 

40 def name(self) -> Name: 

41 """Get the name of this node (last part of the path).""" 

42 if len(self.parts) == 0: 

43 return "" 

44 elif len(self.parts) == 1: 

45 part: Name = self.parts[0] 

46 return part 

47 elif all(isinstance(part, str) for part in self.parts): 

48 return "/".join(quote_part(part) for part in self.parts) 

49 else: 

50 return self 

51 

52 @property 

53 def label(self) -> str: 

54 """Get the label for this node (for display purposes).""" 

55 if len(self.parts) == 0: 

56 return "" 

57 return str(self.parts[-1]) 

58 

59 def drop_root(self, root: Optional["Name"]) -> Optional["NodeKey"]: 

60 """Remove a root prefix from this node key if it matches.""" 

61 if root is None: 

62 return self 

63 root = to_nodekey(root) 

64 n_root_parts = len(root.parts) 

65 if self.is_descendent_of(root): 

66 parts = self.parts[n_root_parts:] 

67 return NodeKey(parts) 

68 else: 

69 return None 

70 

71 def join(self, *others: Name) -> "NodeKey": 

72 """Join this node key with other names to create a new node key.""" 

73 result = self 

74 for other in others: 

75 if other is None: 

76 continue 

77 other = to_nodekey(other) 

78 result = result.join_parts(*other.parts) 

79 return result 

80 

81 def join_parts(self, *parts: Hashable) -> "NodeKey": 

82 """Join this node key with raw parts to create a new node key.""" 

83 if len(parts) == 0: 

84 return self 

85 return NodeKey(self.parts + tuple(parts)) 

86 

87 def __truediv__(self, other: Name) -> "NodeKey": 

88 """Join this node key with other to create a new node key.""" 

89 return self.join(other) 

90 

91 def is_descendent_of(self, other: "NodeKey") -> bool: 

92 """Check if this node key is a descendant of another node key.""" 

93 n_self_parts = len(self.parts) 

94 n_other_parts = len(other.parts) 

95 return n_self_parts > n_other_parts and self.parts[:n_other_parts] == other.parts 

96 

97 @property 

98 def parent(self) -> "NodeKey": 

99 """Get the parent node key.""" 

100 if len(self.parts) == 0: 

101 raise PathNotFoundError() 

102 return NodeKey(self.parts[:-1]) 

103 

104 def prepend(self, nk: "NodeKey") -> "NodeKey": 

105 """Prepend another node key to this one.""" 

106 return nk.join_parts(*self.parts) 

107 

108 def __repr__(self) -> str: 

109 """Return string representation for debugging.""" 

110 path_str = str(self) 

111 quoted_path_str = repr(path_str) 

112 return f"{self.__class__.__name__}({quoted_path_str})" 

113 

114 def __eq__(self, other: object) -> bool: 

115 """Check equality with another NodeKey.""" 

116 if other is None: 

117 return False 

118 if not isinstance(other, NodeKey): 

119 return NotImplemented 

120 return self.parts == other.parts 

121 

122 _ROOT: ClassVar["NodeKey | None"] = None 

123 

124 @classmethod 

125 def root(cls) -> "NodeKey": 

126 """Get the root node key.""" 

127 if cls._ROOT is None: 

128 cls._ROOT = cls(()) 

129 return cls._ROOT 

130 

131 @property 

132 def is_root(self) -> bool: 

133 """Check if this is the root node key.""" 

134 return len(self.parts) == 0 

135 

136 @staticmethod 

137 def common_parent(nodekey1: Name, nodekey2: Name) -> "NodeKey": 

138 """Find the common parent of two node keys.""" 

139 nk1 = to_nodekey(nodekey1) 

140 nk2 = to_nodekey(nodekey2) 

141 parts: list[Hashable] = [] 

142 for p1, p2 in zip(nk1.parts, nk2.parts, strict=False): 

143 if p1 != p2: 

144 break 

145 parts.append(p1) 

146 return NodeKey(tuple(parts)) 

147 

148 def ancestors(self) -> list["NodeKey"]: 

149 """Get all ancestor node keys from root to parent.""" 

150 result = [] 

151 x = self 

152 while True: 

153 result.append(x) 

154 if x.is_root: 

155 break 

156 x = x.parent 

157 return result 

158 

159 

160def names_to_node_keys(names: Name | Names) -> list[NodeKey]: 

161 """Convert names to NodeKey objects.""" 

162 return [to_nodekey(name) for name in as_iterable(names)] 

163 

164 

165def node_keys_to_names(node_keys: Iterable[NodeKey]) -> list[Name]: 

166 """Convert NodeKey objects back to names.""" 

167 return [node_key.name for node_key in node_keys] 

168 

169 

170PART = re.compile(r"([^/]*)/?") 

171 

172 

173def _parse_nodekey(path_str: str, end: int) -> NodeKey: 

174 """Parse a path string into a NodeKey starting from the given position.""" 

175 parts: list[str] = [] 

176 parts_append = parts.append 

177 

178 while path_str[end : end + 1] == "/": 

179 end = end + 1 

180 while True: 

181 nextchar = path_str[end : end + 1] 

182 if nextchar == "": 

183 break 

184 if nextchar == '"': 

185 part, end = json.decoder.scanstring(path_str, end + 1) # type: ignore[attr-defined] 

186 parts_append(part) 

187 nextchar = path_str[end : end + 1] 

188 if nextchar != "" and nextchar != "/": 

189 msg = f"Expected end of string or '/' after quoted part, got {nextchar!r} in {path_str!r}" 

190 raise ValueError(msg) 

191 if nextchar != "": 

192 end = end + 1 

193 else: 

194 chunk = PART.match(path_str, end) 

195 # Defensive: PART = ([^/]*)/? matches the empty string, so match() never returns None here. 

196 if chunk is None: # pragma: no cover 

197 msg = f"Failed to match node key part at position {end} in {path_str!r}" 

198 raise ValueError(msg) 

199 end = chunk.end() 

200 (part,) = chunk.groups() 

201 parts_append(part) 

202 

203 # Defensive: the loop above only breaks once the string is fully consumed, so end always equals its length. 

204 if end != len(path_str): # pragma: no cover 

205 msg = f"Unexpected trailing content at position {end} in {path_str!r}" 

206 raise ValueError(msg) 

207 

208 return NodeKey(tuple(parts)) 

209 

210 

211def parse_nodekey(path_str: str) -> NodeKey: 

212 """Parse a string representation into a NodeKey.""" 

213 return _parse_nodekey(path_str, 0) 

214 

215 

216def to_nodekey(name: Name) -> NodeKey: 

217 """Convert a name to a NodeKey object.""" 

218 if isinstance(name, str): 

219 return parse_nodekey(name) 

220 elif isinstance(name, NodeKey): 

221 return name 

222 elif isinstance(name, object): 

223 return NodeKey((name,)) 

224 else: # pragma: no cover 

225 msg = f"Unexpected error creating node key for name {name}" 

226 raise TypeError(msg) 

227 

228 

229def nodekey_join(*names: Name) -> NodeKey: 

230 """Join multiple names into a single NodeKey.""" 

231 return NodeKey.root().join(*names) 

232 

233 

234def _match_pattern_recursive(pattern: NodeKey, target: NodeKey, p_idx: int, t_idx: int) -> bool: 

235 """Recursively match pattern parts against target parts. 

236 

237 Args: 

238 pattern: The pattern NodeKey to match against 

239 target: The target NodeKey to match 

240 p_idx: Current index in pattern parts 

241 t_idx: Current index in target parts 

242 

243 Returns: 

244 bool: True if pattern matches target, False otherwise 

245 """ 

246 if p_idx == len(pattern.parts) and t_idx == len(target.parts): 

247 return True 

248 if p_idx == len(pattern.parts): 

249 return False 

250 if t_idx == len(target.parts): 

251 return all(p == "**" for p in pattern.parts[p_idx:]) 

252 

253 if pattern.parts[p_idx] == "**": 

254 return _match_pattern_recursive(pattern, target, p_idx + 1, t_idx) or _match_pattern_recursive( 

255 pattern, target, p_idx, t_idx + 1 

256 ) 

257 elif pattern.parts[p_idx] == "*": 

258 return _match_pattern_recursive(pattern, target, p_idx + 1, t_idx + 1) 

259 else: 

260 if pattern.parts[p_idx] == target.parts[t_idx]: 

261 return _match_pattern_recursive(pattern, target, p_idx + 1, t_idx + 1) 

262 return False 

263 

264 

265def is_pattern(nodekey: NodeKey) -> bool: 

266 """Check if a node key contains wildcard patterns.""" 

267 return any(isinstance(part, str) and ("*" in part or "**" in part) for part in nodekey.parts) 

268 

269 

270def match_pattern(pattern: NodeKey, target: NodeKey) -> bool: 

271 """Match a pattern against a target NodeKey. 

272 

273 Supports wildcards: 

274 * - matches exactly one part 

275 ** - matches zero or more parts 

276 

277 Args: 

278 pattern: The pattern to match against 

279 target: The target to match 

280 

281 Returns: 

282 bool: True if pattern matches target, False otherwise 

283 """ 

284 return _match_pattern_recursive(pattern, target, 0, 0)