Coverage for pydantic_evals/pydantic_evals/otel/span_tree.py: 99.74%

1from __future__ import annotations (empty)

2 

3import re (empty)

4from collections.abc import Mapping (empty)

5from datetime import datetime, timedelta, timezone (empty)

6from functools import partial (empty)

7from textwrap import indent (empty)

8from typing import TYPE_CHECKING, Any, Callable (empty)

9 

10from typing_extensions import NotRequired, TypedDict (empty)

11 

12__all__ = 'SpanNode', 'SpanTree', 'SpanQuery', 'as_predicate' (empty)

13 

14if TYPE_CHECKING: # pragma: no cover (empty)

15 # Since opentelemetry isn't a required dependency, don't actually import these at runtime 

16 from opentelemetry.sdk.trace import ReadableSpan 

17 from opentelemetry.trace import SpanContext 

18 from opentelemetry.util.types import AttributeValue 

19 

20 

21class SpanNode: (empty)

22 """A node in the span tree; provides references to parents/children for easy traversal and queries.""" 

23 

24 def __init__(self, span: ReadableSpan): (empty)

25 self._span = span (empty)

26 # If a span has no context, it's going to cause problems. We may need to add improved handling of this scenario. 

27 assert self._span.context is not None, f'{span=} has no context' (empty)

28 

29 self.parent: SpanNode | None = None (empty)

30 self.children_by_id: dict[int, SpanNode] = {} # note: we rely on insertion order to determine child order (empty)

31 

32 @property (empty)

33 def children(self) -> list[SpanNode]: (empty)

34 return list(self.children_by_id.values()) (empty)

35 

36 @property (empty)

37 def descendants(self) -> list[SpanNode]: (empty)

38 """Return all descendants of this node in DFS order.""" 

39 descendants: list[SpanNode] = [] (empty)

40 stack = list(self.children) (empty)

41 while stack: (empty)

42 node = stack.pop() (empty)

43 descendants.append(node) (empty)

44 stack.extend(node.children) (empty)

45 return descendants (empty)

46 

47 @property (empty)

48 def context(self) -> SpanContext: (empty)

49 """Return the SpanContext of the wrapped span.""" 

50 assert self._span.context is not None (empty)

51 return self._span.context (empty)

52 

53 @property (empty)

54 def parent_context(self) -> SpanContext | None: (empty)

55 """Return the SpanContext of the parent of the wrapped span.""" 

56 return self._span.parent (empty)

57 

58 @property (empty)

59 def span_id(self) -> int: (empty)

60 """Return the integer span_id from the SpanContext.""" 

61 return self.context.span_id (empty)

62 

63 @property (empty)

64 def trace_id(self) -> int: (empty)

65 """Return the integer trace_id from the SpanContext.""" 

66 return self.context.trace_id (empty)

67 

68 @property (empty)

69 def name(self) -> str: (empty)

70 """Convenience for the span's name.""" 

71 return self._span.name (empty)

72 

73 @property (empty)

74 def start_timestamp(self) -> datetime: (empty)

75 """Return the span's start time as a UTC datetime, or None if not set.""" 

76 assert self._span.start_time is not None (empty)

77 return datetime.fromtimestamp(self._span.start_time / 1e9, tz=timezone.utc) (empty)

78 

79 @property (empty)

80 def end_timestamp(self) -> datetime: (empty)

81 """Return the span's end time as a UTC datetime, or None if not set.""" 

82 assert self._span.end_time is not None (empty)

83 return datetime.fromtimestamp(self._span.end_time / 1e9, tz=timezone.utc) (empty)

84 

85 @property (empty)

86 def duration(self) -> timedelta: (empty)

87 """Return the span's duration as a timedelta, or None if start/end not set.""" 

88 return self.end_timestamp - self.start_timestamp (empty)

89 

90 @property (empty)

91 def attributes(self) -> Mapping[str, AttributeValue]: (empty)

92 # Note: It would be nice to expose the non-JSON-serialized versions of (logfire-recorded) attributes with 

93 # nesting etc. This just exposes the JSON-serialized version, but doing more would be difficult. 

94 return self._span.attributes or {} (empty)

95 

96 def add_child(self, child: SpanNode) -> None: (empty)

97 """Attach a child node to this node's list of children.""" 

98 self.children_by_id[child.span_id] = child (empty)

99 child.parent = self (empty)

100 

101 # ------------------------------------------------------------------------- 

102 # Child queries 

103 # ------------------------------------------------------------------------- 

104 def find_children(self, predicate: Callable[[SpanNode], bool]) -> list[SpanNode]: (empty)

105 """Return all immediate children that satisfy the given predicate.""" 

106 return [child for child in self.children if predicate(child)] (empty)

107 

108 def first_child(self, predicate: Callable[[SpanNode], bool]) -> SpanNode | None: (empty)

109 """Return the first immediate child that satisfies the given predicate, or None if none match.""" 

110 for child in self.children: (empty)

111 if predicate(child): (empty)

112 return child (empty)

113 return None (empty)

114 

115 def any_child(self, predicate: Callable[[SpanNode], bool]) -> bool: (empty)

116 """Returns True if there is at least one child that satisfies the predicate.""" 

117 return self.first_child(predicate) is not None (empty)

118 

119 # ------------------------------------------------------------------------- 

120 # Descendant queries (DFS) 

121 # ------------------------------------------------------------------------- 

122 def find_descendants(self, predicate: Callable[[SpanNode], bool]) -> list[SpanNode]: (empty)

123 """Return all descendant nodes that satisfy the given predicate in DFS order.""" 

124 found: list[SpanNode] = [] (empty)

125 stack = list(self.children) (empty)

126 while stack: (empty)

127 node = stack.pop() (empty)

128 if predicate(node): (empty)

129 found.append(node) (empty)

130 stack.extend(node.children) (empty)

131 return found (empty)

132 

133 def first_descendant(self, predicate: Callable[[SpanNode], bool]) -> SpanNode | None: (empty)

134 """DFS: Return the first descendant (in DFS order) that satisfies the given predicate, or `None` if none match.""" 

135 stack = list(self.children) (empty)

136 while stack: (empty)

137 node = stack.pop() (empty)

138 if predicate(node): (empty)

139 return node (empty)

140 stack.extend(node.children) (empty)

141 return None (empty)

142 

143 def any_descendant(self, predicate: Callable[[SpanNode], bool]) -> bool: (empty)

144 """Returns `True` if there is at least one descendant that satisfies the predicate.""" 

145 return self.first_descendant(predicate) is not None (empty)

146 

147 # ------------------------------------------------------------------------- 

148 # Ancestor queries (DFS "up" the chain) 

149 # ------------------------------------------------------------------------- 

150 def find_ancestors(self, predicate: Callable[[SpanNode], bool]) -> list[SpanNode]: (empty)

151 """Return all ancestors that satisfy the given predicate.""" 

152 found: list[SpanNode] = [] (empty)

153 node = self.parent (empty)

154 while node: (empty)

155 if predicate(node): 155 ↛ 157line 155 didn't jump to line 157 because the condition on line 155 was always true(empty)

156 found.append(node) (empty)

157 node = node.parent (empty)

158 return found (empty)

159 

160 def first_ancestor(self, predicate: Callable[[SpanNode], bool]) -> SpanNode | None: (empty)

161 """Return the closest ancestor that satisfies the given predicate, or `None` if none match.""" 

162 node = self.parent (empty)

163 while node: (empty)

164 if predicate(node): (empty)

165 return node (empty)

166 node = node.parent (empty)

167 return None (empty)

168 

169 def any_ancestor(self, predicate: Callable[[SpanNode], bool]) -> bool: (empty)

170 """Returns True if any ancestor satisfies the predicate.""" 

171 return self.first_ancestor(predicate) is not None (empty)

172 

173 # ------------------------------------------------------------------------- 

174 # Matching convenience 

175 # ------------------------------------------------------------------------- 

176 def matches(self, name: str | None = None, attributes: dict[str, Any] | None = None) -> bool: (empty)

177 """A convenience method to see if this node's span matches certain conditions. 

178 

179 - name: exact match for the Span name 

180 - attributes: dict of key->value; must match exactly. 

181 """ 

182 if name is not None and self.name != name: (empty)

183 return False (empty)

184 if attributes: (empty)

185 span_attributes = self._span.attributes or {} (empty)

186 for attr_key, attr_val in attributes.items(): (empty)

187 if span_attributes.get(attr_key) != attr_val: (empty)

188 return False (empty)

189 return True (empty)

190 

191 # ------------------------------------------------------------------------- 

192 # String representation 

193 # ------------------------------------------------------------------------- 

194 def repr_xml( (empty)

195 self, 

196 include_children: bool = True, 

197 include_span_id: bool = False, 

198 include_trace_id: bool = False, 

199 include_start_timestamp: bool = False, 

200 include_duration: bool = False, 

201 ) -> str: 

202 """Return an XML-like string representation of the node. 

203 

204 Optionally includes children, span_id, trace_id, start_timestamp, and duration. 

205 """ 

206 first_line_parts = [f'<SpanNode name={self.name!r}'] (empty)

207 if include_span_id: (empty)

208 first_line_parts.append(f'span_id={self.span_id:016x}') (empty)

209 if include_trace_id: (empty)

210 first_line_parts.append(f'trace_id={self.trace_id:032x}') (empty)

211 if include_start_timestamp: (empty)

212 first_line_parts.append(f'start_timestamp={self.start_timestamp.isoformat()!r}') (empty)

213 if include_duration: (empty)

214 first_line_parts.append(f"duration='{self.duration}'") (empty)

215 

216 extra_lines: list[str] = [] (empty)

217 if include_children and self.children: (empty)

218 first_line_parts.append('>') (empty)

219 for child in self.children: (empty)

220 extra_lines.append( (empty)

221 indent( 

222 child.repr_xml( 

223 include_children=include_children, 

224 include_span_id=include_span_id, 

225 include_trace_id=include_trace_id, 

226 include_start_timestamp=include_start_timestamp, 

227 include_duration=include_duration, 

228 ), 

229 ' ', 

230 ) 

231 ) 

232 extra_lines.append('</SpanNode>') (empty)

233 else: 

234 if self.children: (empty)

235 first_line_parts.append('children=...') (empty)

236 first_line_parts.append('/>') (empty)

237 return '\n'.join([' '.join(first_line_parts), *extra_lines]) (empty)

238 

239 def __str__(self) -> str: (empty)

240 if self.children: (empty)

241 return f'<SpanNode name={self.name!r} span_id={self.span_id:016x}>...</SpanNode>' (empty)

242 else: 

243 return f'<SpanNode name={self.name!r} span_id={self.span_id:016x} />' (empty)

244 

245 def __repr__(self) -> str: (empty)

246 return self.repr_xml() (empty)

247 

248 

249class SpanTree: (empty)

250 """A container that builds a hierarchy of SpanNode objects from a list of finished spans. 

251 

252 You can then search or iterate the tree to make your assertions (using DFS for traversal). 

253 """ 

254 

255 def __init__(self, spans: list[ReadableSpan] | None = None): (empty)

256 self.nodes_by_id: dict[int, SpanNode] = {} (empty)

257 self.roots: list[SpanNode] = [] (empty)

258 if spans: # pragma: no cover (empty)

259 self.add_spans(spans) 

260 

261 def add_spans(self, spans: list[ReadableSpan]) -> None: (empty)

262 """Add a list of spans to the tree, rebuilding the tree structure.""" 

263 for span in spans: (empty)

264 node = SpanNode(span) (empty)

265 self.nodes_by_id[node.span_id] = node (empty)

266 self._rebuild_tree() (empty)

267 

268 def _rebuild_tree(self): (empty)

269 # Ensure spans are ordered by start_timestamp so that roots and children end up in the right order 

270 nodes = list(self.nodes_by_id.values()) (empty)

271 nodes.sort(key=lambda node: node.start_timestamp or datetime.min) (empty)

272 self.nodes_by_id = {node.span_id: node for node in nodes} (empty)

273 

274 # Build the parent/child relationships 

275 for node in self.nodes_by_id.values(): (empty)

276 parent_ctx = node.parent_context (empty)

277 if parent_ctx is not None: (empty)

278 parent_node = self.nodes_by_id.get(parent_ctx.span_id) (empty)

279 if parent_node is not None: (empty)

280 parent_node.add_child(node) (empty)

281 

282 # Determine the roots 

283 # A node is a "root" if its parent is None or if its parent's span_id is not in the current set of spans. 

284 self.roots = [] (empty)

285 for node in self.nodes_by_id.values(): (empty)

286 parent_ctx = node.parent_context (empty)

287 if parent_ctx is None or parent_ctx.span_id not in self.nodes_by_id: (empty)

288 self.roots.append(node) (empty)

289 

290 def flattened(self) -> list[SpanNode]: (empty)

291 """Return a list of all nodes in the tree.""" 

292 return list(self.nodes_by_id.values()) (empty)

293 

294 def find_all(self, predicate: Callable[[SpanNode], bool]) -> list[SpanNode]: (empty)

295 """Find all nodes in the entire tree that match the predicate, scanning from each root in DFS order.""" 

296 result: list[SpanNode] = [] (empty)

297 stack = self.roots[:] (empty)

298 while stack: (empty)

299 node = stack.pop() (empty)

300 if predicate(node): (empty)

301 result.append(node) (empty)

302 stack.extend(node.children) (empty)

303 return result (empty)

304 

305 def find_first(self, predicate: Callable[[SpanNode], bool]) -> SpanNode | None: (empty)

306 """Find the first node that matches a predicate, scanning from each root in DFS order. Returns `None` if not found.""" 

307 stack = self.roots[:] (empty)

308 while stack: (empty)

309 node = stack.pop() (empty)

310 if predicate(node): (empty)

311 return node (empty)

312 stack.extend(node.children) (empty)

313 return None (empty)

314 

315 def any(self, predicate: Callable[[SpanNode], bool]) -> bool: (empty)

316 """Returns True if any node in the tree matches the predicate.""" 

317 return self.find_first(predicate) is not None (empty)

318 

319 def __str__(self): (empty)

320 return f'<SpanTree num_roots={len(self.roots)} total_spans={len(self.nodes_by_id)} />' (empty)

321 

322 def repr_xml( (empty)

323 self, 

324 include_children: bool = True, 

325 include_span_id: bool = False, 

326 include_trace_id: bool = False, 

327 include_start_timestamp: bool = False, 

328 include_duration: bool = False, 

329 ) -> str: 

330 """Return an XML-like string representation of the tree, optionally including children, span_id, trace_id, duration, and timestamps.""" 

331 if not self.roots: (empty)

332 return '<SpanTree />' (empty)

333 repr_parts = [ (empty)

334 '<SpanTree>', 

335 *[ 

336 indent( 

337 root.repr_xml( 

338 include_children=include_children, 

339 include_span_id=include_span_id, 

340 include_trace_id=include_trace_id, 

341 include_start_timestamp=include_start_timestamp, 

342 include_duration=include_duration, 

343 ), 

344 ' ', 

345 ) 

346 for root in self.roots 

347 ], 

348 '</SpanTree>', 

349 ] 

350 return '\n'.join(repr_parts) (empty)

351 

352 def __repr__(self): (empty)

353 return self.repr_xml() (empty)

354 

355 

356class SpanQuery(TypedDict): (empty)

357 """A serializable query for filtering SpanNodes based on various conditions. 

358 

359 All fields are optional and combined with AND logic by default. 

360 

361 Due to the presence of `__calL__`, a `SpanQuery` can be used as a predicate in `SpanTree.find_first`, etc. 

362 """ 

363 

364 # Individual span conditions 

365 ## Name conditions 

366 name_equals: NotRequired[str] (empty)

367 name_contains: NotRequired[str] (empty)

368 name_matches_regex: NotRequired[str] # regex pattern (empty)

369 

370 ## Attribute conditions 

371 has_attributes: NotRequired[dict[str, Any]] (empty)

372 has_attribute_keys: NotRequired[list[str]] (empty)

373 

374 ## Timing conditions 

375 min_duration: NotRequired[timedelta | float] (empty)

376 max_duration: NotRequired[timedelta | float] (empty)

377 

378 # Logical combinations of conditions 

379 not_: NotRequired[SpanQuery] (empty)

380 and_: NotRequired[list[SpanQuery]] (empty)

381 or_: NotRequired[list[SpanQuery]] (empty)

382 

383 # Descendant conditions 

384 some_child_has: NotRequired[SpanQuery] (empty)

385 all_children_have: NotRequired[SpanQuery] (empty)

386 no_child_has: NotRequired[SpanQuery] (empty)

387 min_child_count: NotRequired[int] (empty)

388 max_child_count: NotRequired[int] (empty)

389 

390 some_descendant_has: NotRequired[SpanQuery] (empty)

391 all_descendants_have: NotRequired[SpanQuery] (empty)

392 no_descendant_has: NotRequired[SpanQuery] (empty)

393 

394 

395def as_predicate(query: SpanQuery) -> Callable[[SpanNode], bool]: (empty)

396 """Convert a SpanQuery into a callable predicate that can be used in SpanTree.find_first, etc.""" 

397 return partial(matches, query) (empty)

398 

399 

400def matches(query: SpanQuery, span: SpanNode) -> bool: # noqa C901 (empty)

401 """Check if the span matches the query conditions.""" 

402 # Logical combinations 

403 if or_ := query.get('or_'): (empty)

404 if len(query) > 1: (empty)

405 raise ValueError("Cannot combine 'or_' conditions with other conditions at the same level") (empty)

406 return any(matches(q, span) for q in or_) (empty)

407 if not_ := query.get('not_'): (empty)

408 if matches(not_, span): (empty)

409 return False (empty)

410 if and_ := query.get('and_'): (empty)

411 results = [matches(q, span) for q in and_] (empty)

412 if not all(results): (empty)

413 return False (empty)

414 # At this point, all existing ANDs and no existing ORs have passed, so it comes down to this condition 

415 

416 # Name conditions 

417 if (name_equals := query.get('name_equals')) and span.name != name_equals: (empty)

418 return False (empty)

419 if (name_contains := query.get('name_contains')) and name_contains not in span.name: (empty)

420 return False (empty)

421 if (name_matches_regex := query.get('name_matches_regex')) and not re.match(name_matches_regex, span.name): (empty)

422 return False (empty)

423 

424 # Attribute conditions 

425 if (has_attributes := query.get('has_attributes')) and not all( (empty)

426 span.attributes.get(key) == value for key, value in has_attributes.items() 

427 ): 

428 return False (empty)

429 if (has_attributes_keys := query.get('has_attribute_keys')) and not all( (empty)

430 key in span.attributes for key in has_attributes_keys 

431 ): 

432 return False (empty)

433 

434 # Timing conditions 

435 if (min_duration := query.get('min_duration')) is not None and span.duration is not None: (empty)

436 if not isinstance(min_duration, timedelta): (empty)

437 min_duration = timedelta(seconds=min_duration) (empty)

438 if span.duration < min_duration: (empty)

439 return False (empty)

440 if (max_duration := query.get('max_duration')) is not None and span.duration is not None: (empty)

441 if not isinstance(max_duration, timedelta): (empty)

442 max_duration = timedelta(seconds=max_duration) (empty)

443 if span.duration > max_duration: (empty)

444 return False (empty)

445 

446 # Children conditions 

447 if (min_child_count := query.get('min_child_count')) and len(span.children) < min_child_count: (empty)

448 return False (empty)

449 if (max_child_count := query.get('max_child_count')) and len(span.children) > max_child_count: (empty)

450 return False (empty)

451 if (some_child_has := query.get('some_child_has')) and not any( (empty)

452 matches(some_child_has, child) for child in span.children 

453 ): 

454 return False (empty)

455 if (all_children_have := query.get('all_children_have')) and not all( (empty)

456 matches(all_children_have, child) for child in span.children 

457 ): 

458 return False (empty)

459 if (no_child_has := query.get('no_child_has')) and any(matches(no_child_has, child) for child in span.children): (empty)

460 return False (empty)

461 

462 # Descendant conditions 

463 if (some_descendant_has := query.get('some_descendant_has')) and not any( (empty)

464 matches(some_descendant_has, child) for child in span.descendants 

465 ): 

466 return False (empty)

467 if (all_descendants_have := query.get('all_descendants_have')) and not all( (empty)

468 matches(all_descendants_have, child) for child in span.descendants 

469 ): 

470 return False (empty)

471 if (no_descendant_has := query.get('no_descendant_has')) and any( (empty)

472 matches(no_descendant_has, child) for child in span.descendants 

473 ): 

474 return False (empty)

475 

476 return True (empty)