1use std::collections::{BTreeMap, BTreeSet};
12
13use ruff_python_ast::{
14 BoolOp, Comprehension, Expr, Stmt, UnaryOp,
15 helpers::is_docstring_stmt,
16 visitor::{Visitor, walk_comprehension, walk_expr, walk_stmt},
17};
18use ruff_python_parser::parse_module;
19use ruff_text_size::{Ranged, TextRange};
20use serde::{Deserialize, Serialize};
21use sha2::{Digest, Sha256};
22
23use crate::{
24 coverage_analysis::PointKind,
25 coverage_report::{
26 BranchAlternativeMeta, BranchMeta, CoverageManifest, DecisionMeta, PointMeta,
27 },
28};
29
30pub const PYTHON_PROBE_PLAN_VERSION: u32 = 1;
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub enum PythonInstrumenterError {
34 SourceTooLarge,
35 Parse(String),
36 InvalidRange,
37}
38
39impl std::fmt::Display for PythonInstrumenterError {
40 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41 match self {
42 Self::SourceTooLarge => write!(formatter, "Python source exceeds the parser range"),
43 Self::Parse(error) => write!(formatter, "Python parse failed: {error}"),
44 Self::InvalidRange => write!(formatter, "Python parser returned an invalid range"),
45 }
46 }
47}
48
49impl std::error::Error for PythonInstrumenterError {}
50
51#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
55#[serde(from = "[[usize; 2]; 2]", into = "[[usize; 2]; 2]")]
56pub struct PlanSpan {
57 pub start: [usize; 2],
58 pub end: [usize; 2],
59}
60
61impl From<[[usize; 2]; 2]> for PlanSpan {
62 fn from(value: [[usize; 2]; 2]) -> Self {
63 Self {
64 start: value[0],
65 end: value[1],
66 }
67 }
68}
69
70impl From<PlanSpan> for [[usize; 2]; 2] {
71 fn from(value: PlanSpan) -> Self {
72 [value.start, value.end]
73 }
74}
75
76#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
77#[serde(rename_all = "camelCase", deny_unknown_fields)]
78pub struct StatementPlan {
79 pub id: String,
80 pub start: [usize; 2],
83}
84
85#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
86#[serde(rename_all = "camelCase", deny_unknown_fields)]
87pub struct FunctionPlan {
88 pub id: String,
89 pub name: String,
90 pub span: PlanSpan,
92 #[serde(default, skip_serializing_if = "Option::is_none")]
96 pub first: Option<String>,
97}
98
99#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
100#[serde(rename_all = "camelCase", deny_unknown_fields)]
101pub struct HandlerPlan {
102 pub id: String,
103 pub bare: bool,
104 pub missed: String,
105 pub selected: String,
106}
107
108#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
109#[serde(rename_all = "camelCase", deny_unknown_fields)]
110pub struct TryPlan {
111 pub id: String,
112 pub body: PlanSpan,
113 #[serde(skip_serializing_if = "Option::is_none")]
114 pub finalbody: Option<PlanSpan>,
115 pub handlers: Vec<HandlerPlan>,
116 pub success: String,
117 pub raised: String,
118}
119
120#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
121#[serde(rename_all = "camelCase", deny_unknown_fields)]
122pub struct ConditionPlan {
123 pub span: PlanSpan,
124 pub not: usize,
127}
128
129#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
134#[serde(untagged)]
135pub enum ConditionTree {
136 Leaf(usize),
137 Node {
138 op: String,
139 items: Vec<ConditionTree>,
140 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
141 negate: bool,
142 },
143}
144
145#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
146#[serde(rename_all = "camelCase", deny_unknown_fields)]
147pub struct DecisionPlan {
148 pub id: String,
149 pub kind: String,
150 pub span: PlanSpan,
151 pub conditions: Vec<ConditionPlan>,
152 pub outcome_true: String,
153 pub outcome_false: String,
154 #[serde(default, skip_serializing_if = "Option::is_none")]
159 pub when_true: Option<String>,
160 #[serde(default, skip_serializing_if = "Option::is_none")]
161 pub when_false: Option<String>,
162 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
165 pub value: bool,
166}
167
168#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
169#[serde(rename_all = "camelCase", deny_unknown_fields)]
170pub struct LoopPlan {
171 pub id: String,
172 pub iter: PlanSpan,
173 pub zero: String,
174 pub entered: String,
175 #[serde(default, skip_serializing_if = "Option::is_none")]
178 pub first: Option<String>,
179}
180
181#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
182#[serde(rename_all = "camelCase", deny_unknown_fields)]
183pub struct LogicalPlan {
184 pub id: String,
185 pub boolop: PlanSpan,
186 pub operand: usize,
189 #[serde(skip_serializing_if = "Option::is_none")]
193 pub decision: Option<String>,
194 #[serde(skip_serializing_if = "Option::is_none")]
195 pub previous_leaves: Option<Vec<usize>>,
196 #[serde(skip_serializing_if = "Option::is_none")]
197 pub operand_leaves: Option<Vec<usize>>,
198 pub short_circuit: String,
199 pub evaluated: String,
200}
201
202#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
203#[serde(rename_all = "camelCase", deny_unknown_fields)]
204pub struct MatchCasePlan {
205 pub id: String,
206 pub span: PlanSpan,
207 pub irrefutable: bool,
208 pub missed: String,
209 pub selected: String,
210}
211
212#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
213#[serde(rename_all = "camelCase", deny_unknown_fields)]
214pub struct MatchNoCasePlan {
215 pub id: String,
216 pub matched: String,
217 pub unmatched: String,
218}
219
220#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
221#[serde(rename_all = "camelCase", deny_unknown_fields)]
222pub struct MatchPlan {
223 pub span: PlanSpan,
224 pub cases: Vec<MatchCasePlan>,
225 #[serde(skip_serializing_if = "Option::is_none")]
226 pub no_case: Option<MatchNoCasePlan>,
227}
228
229#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
230#[serde(rename_all = "camelCase", deny_unknown_fields)]
231pub struct PythonFilePlan {
232 pub statements: Vec<StatementPlan>,
233 pub functions: Vec<FunctionPlan>,
234 pub decisions: Vec<DecisionPlan>,
235 pub loops: Vec<LoopPlan>,
236 pub logical: Vec<LogicalPlan>,
237 pub matches: Vec<MatchPlan>,
238 #[serde(default, skip_serializing_if = "Vec::is_empty")]
239 pub tries: Vec<TryPlan>,
240}
241
242#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
243#[serde(rename_all = "camelCase", deny_unknown_fields)]
244pub struct PythonProbePlan {
245 pub version: u32,
246 pub root: String,
247 pub files: BTreeMap<String, PythonFilePlan>,
248 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
255 pub assertion_sites: BTreeMap<String, Vec<usize>>,
256}
257
258#[derive(Debug, Clone, PartialEq)]
259pub struct PythonFileObligations {
260 pub manifest: CoverageManifest,
261 pub plan: PythonFilePlan,
262}
263
264struct SourceLocations<'a> {
265 source: &'a str,
266 line_starts: Vec<usize>,
267}
268
269impl<'a> SourceLocations<'a> {
270 fn new(source: &'a str) -> Self {
271 let mut line_starts = vec![0];
272 line_starts.extend(
273 source
274 .bytes()
275 .enumerate()
276 .filter_map(|(index, byte)| (byte == b'\n').then_some(index + 1)),
277 );
278 Self {
279 source,
280 line_starts,
281 }
282 }
283
284 fn range(&self, range: TextRange) -> Result<(usize, usize), PythonInstrumenterError> {
285 let start = range.start().to_usize();
286 let end = range.end().to_usize();
287 if start > end
288 || end > self.source.len()
289 || !self.source.is_char_boundary(start)
290 || !self.source.is_char_boundary(end)
291 {
292 return Err(PythonInstrumenterError::InvalidRange);
293 }
294 Ok((start, end))
295 }
296
297 fn line_column(&self, offset: usize) -> (usize, usize) {
298 let line_index = self.line_starts.partition_point(|start| *start <= offset) - 1;
299 (line_index + 1, offset - self.line_starts[line_index])
300 }
301
302 fn span(&self, range: TextRange) -> Result<PlanSpan, PythonInstrumenterError> {
303 let (start, end) = self.range(range)?;
304 let (start_line, start_column) = self.line_column(start);
305 let (end_line, end_column) = self.line_column(end);
306 Ok(PlanSpan {
307 start: [start_line, start_column],
308 end: [end_line, end_column],
309 })
310 }
311
312 fn text(&self, range: TextRange) -> Result<String, PythonInstrumenterError> {
313 let (start, end) = self.range(range)?;
314 Ok(self.source[start..end].trim().to_owned())
315 }
316}
317
318fn stable_id(file: &str, kind: &str, range: TextRange, suffix: &str) -> String {
319 let mut hash = Sha256::new();
320 let start = range.start().to_usize().to_string();
321 let end = range.end().to_usize().to_string();
322 for value in [file, kind, &start, &end, suffix] {
323 hash.update(value.as_bytes());
324 hash.update([0]);
325 }
326 let digest = hash.finalize();
327 let mut encoded = String::with_capacity(24);
328 for byte in &digest[..12] {
329 use std::fmt::Write as _;
330 write!(&mut encoded, "{byte:02x}").expect("writing to a string cannot fail");
331 }
332 format!("py:{kind}:{encoded}")
333}
334
335fn is_unobservable_statement(statement: &Stmt, first_in_body: bool) -> bool {
339 matches!(statement, Stmt::Global(_) | Stmt::Nonlocal(_))
340 || (first_in_body && is_docstring_stmt(statement))
341}
342
343type DecisionBoolOp = (String, Vec<Vec<usize>>);
346
347type RegroupedBoolOp = (String, Vec<Vec<usize>>, PlanSpan);
350
351type TreeBoolOps = Vec<(TextRange, Vec<Vec<usize>>)>;
354
355struct PythonObligationCollector<'a> {
356 file: &'a str,
357 locations: SourceLocations<'a>,
358 manifest: CoverageManifest,
359 plan: PythonFilePlan,
360 point_ids: BTreeSet<String>,
361 decision_ids: BTreeSet<String>,
362 branch_ids: BTreeSet<String>,
363 decision_boolops: BTreeMap<(usize, usize), DecisionBoolOp>,
369 regrouped: BTreeMap<(usize, usize), RegroupedBoolOp>,
373 error: Option<PythonInstrumenterError>,
374}
375
376impl<'a> PythonObligationCollector<'a> {
377 fn new(file: &'a str, source: &'a str) -> Self {
378 Self {
379 file,
380 locations: SourceLocations::new(source),
381 manifest: CoverageManifest {
382 unmeasured: Vec::new(),
383 decisions: Vec::new(),
384 points: Vec::new(),
385 branches: Vec::new(),
386 limitations: Vec::new(),
387 scope: None,
388 assertion_sites: Vec::new(),
389 },
390 plan: PythonFilePlan::default(),
391 point_ids: BTreeSet::new(),
392 decision_ids: BTreeSet::new(),
393 branch_ids: BTreeSet::new(),
394 decision_boolops: BTreeMap::new(),
395 regrouped: BTreeMap::new(),
396 error: None,
397 }
398 }
399
400 fn fail<T>(&mut self, error: PythonInstrumenterError) -> Option<T> {
401 self.error.get_or_insert(error);
402 None
403 }
404
405 fn location_source(&mut self, range: TextRange) -> Option<(usize, usize, String)> {
406 let result = self.locations.range(range).map(|(start, _)| {
407 let (line, column) = self.locations.line_column(start);
408 (line, column, self.locations.text(range))
409 });
410 match result {
411 Ok((line, column, Ok(source))) => Some((line, column, source)),
412 Ok((_, _, Err(error))) | Err(error) => self.fail(error),
413 }
414 }
415
416 fn span(&mut self, range: TextRange) -> Option<PlanSpan> {
417 match self.locations.span(range) {
418 Ok(span) => Some(span),
419 Err(error) => self.fail(error),
420 }
421 }
422
423 fn push_point(&mut self, id: &str, range: TextRange, kind: PointKind, label: Option<String>) {
424 let Some((line, column, source)) = self.location_source(range) else {
425 return;
426 };
427 self.manifest.points.push(PointMeta {
428 id: id.into(),
429 kind,
430 file: self.file.into(),
431 line,
432 column,
433 source,
434 label,
435 });
436 }
437
438 fn statement(&mut self, statement: &Stmt) {
439 let range = statement.range();
440 let id = stable_id(self.file, "statement", range, "");
441 if !self.point_ids.insert(id.clone()) {
442 return;
443 }
444 self.push_point(&id, range, PointKind::Statement, None);
445 let Some(span) = self.span(range) else {
446 return;
447 };
448 self.plan.statements.push(StatementPlan {
449 id,
450 start: span.start,
451 });
452 }
453
454 fn function(&mut self, range: TextRange, name: &str) {
455 let id = stable_id(self.file, "function", range, name);
456 if !self.point_ids.insert(id.clone()) {
457 return;
458 }
459 self.push_point(&id, range, PointKind::Function, Some(name.to_owned()));
460 let Some(span) = self.span(range) else {
461 return;
462 };
463 self.plan.functions.push(FunctionPlan {
464 id,
465 name: name.to_owned(),
466 span,
467 first: None,
468 });
469 }
470
471 fn strip_not(expr: &Expr) -> (&Expr, usize) {
472 let mut current = expr;
473 let mut depth = 0;
474 while let Expr::UnaryOp(unary) = current {
475 if unary.op != UnaryOp::Not {
476 break;
477 }
478 depth += 1;
479 current = &unary.operand;
480 }
481 (current, depth)
482 }
483
484 fn tree(
490 &mut self,
491 expr: &Expr,
492 leaves: &mut Vec<ConditionPlan>,
493 boolops: &mut Vec<(TextRange, Vec<Vec<usize>>)>,
494 ) -> Option<ConditionTree> {
495 let (operand, not) = Self::strip_not(expr);
496 if let Expr::BoolOp(boolean) = operand {
497 let mut items = Vec::with_capacity(boolean.values.len());
498 let mut groups = Vec::with_capacity(boolean.values.len());
499 for value in &boolean.values {
500 let first = leaves.len();
501 items.push(self.tree(value, leaves, boolops)?);
502 groups.push((first..leaves.len()).collect());
503 }
504 boolops.push((boolean.range, groups));
505 return Some(ConditionTree::Node {
506 op: match boolean.op {
507 BoolOp::And => "and".into(),
508 BoolOp::Or => "or".into(),
509 },
510 items,
511 negate: not % 2 == 1,
512 });
513 }
514 let span = self.span(operand.range())?;
515 leaves.push(ConditionPlan { span, not });
516 Some(ConditionTree::Leaf(leaves.len() - 1))
517 }
518
519 fn decision(&mut self, test: &Expr, kind: &str) {
520 self.decision_over(test.range(), kind, false, |this, leaves, boolops| {
521 this.tree(test, leaves, boolops).map(|_| ())
522 });
523 }
524
525 fn decision_over(
528 &mut self,
529 range: TextRange,
530 kind: &str,
531 value: bool,
532 build: impl FnOnce(&mut Self, &mut Vec<ConditionPlan>, &mut TreeBoolOps) -> Option<()>,
533 ) -> Option<(String, TreeBoolOps)> {
534 let id = stable_id(self.file, "decision", range, kind);
535 if !self.decision_ids.insert(id.clone()) {
536 return None;
537 }
538 let (line, column, source) = self.location_source(range)?;
539 let mut leaves = Vec::new();
540 let mut boolops = Vec::new();
541 build(self, &mut leaves, &mut boolops)?;
542 let tree = boolops.clone();
543 let mut conditions = Vec::with_capacity(leaves.len());
544 for leaf in &leaves {
545 let leaf_range = TextRange::new(
546 self.locations.line_starts[leaf.span.start[0] - 1]
547 .saturating_add(leaf.span.start[1])
548 .try_into()
549 .unwrap_or_default(),
550 self.locations.line_starts[leaf.span.end[0] - 1]
551 .saturating_add(leaf.span.end[1])
552 .try_into()
553 .unwrap_or_default(),
554 );
555 match self.locations.text(leaf_range) {
556 Ok(text) => conditions.push(if leaf.not % 2 == 1 {
557 format!("not {text}")
558 } else {
559 text
560 }),
561 Err(error) => {
562 return self.fail(error);
563 }
564 }
565 }
566 for (boolop_range, groups) in boolops {
567 let (start, end) = match self.locations.range(boolop_range) {
568 Ok(bounds) => bounds,
569 Err(error) => {
570 return self.fail(error);
571 }
572 };
573 self.decision_boolops
574 .insert((start, end), (id.clone(), groups));
575 }
576 self.manifest.decisions.push(DecisionMeta {
577 id: id.clone(),
578 file: self.file.into(),
579 line,
580 column,
581 source: source.clone(),
582 conditions,
583 kind: kind.into(),
584 });
585 let outcome_id = format!("{id}:outcome");
586 self.branch_with_id(
587 outcome_id.clone(),
588 range,
589 kind,
590 source,
591 [("true", "true"), ("false", "false")],
592 );
593 let span = self.span(range)?;
594 self.plan.decisions.push(DecisionPlan {
595 id: id.clone(),
596 kind: kind.into(),
597 span,
598 conditions: leaves,
599 outcome_true: format!("{outcome_id}:true"),
600 outcome_false: format!("{outcome_id}:false"),
601 when_true: None,
602 when_false: None,
603 value,
604 });
605 Some((id, tree))
606 }
607
608 fn value_decision(&mut self, expression: &Expr) {
614 let bounds = |this: &mut Self, range: TextRange| this.locations.range(range).ok();
615 match expression {
616 Expr::BoolOp(boolean) => {
617 let Some(own) = bounds(self, boolean.range) else {
618 return;
619 };
620 if self.decision_boolops.contains_key(&own) {
621 return;
622 }
623 let values = &boolean.values;
624 if values.len() >= 3 {
625 let left = &values[..values.len() - 1];
630 let range =
631 TextRange::new(left[0].range().start(), left[left.len() - 1].range().end());
632 let Some(span) = self.span(range) else {
633 return;
634 };
635 if let Some((id, tree)) =
636 self.decision_over(range, "value", true, |this, leaves, boolops| {
637 let mut groups = Vec::with_capacity(left.len());
638 for value in left {
639 let first = leaves.len();
640 this.tree(value, leaves, boolops)?;
641 groups.push((first..leaves.len()).collect());
642 }
643 boolops.push((range, groups));
644 Some(())
645 })
646 && let Some((_, groups)) =
647 tree.into_iter().find(|(boolop, _)| *boolop == range)
648 {
649 self.regrouped.insert(own, (id, groups, span));
650 }
651 } else if Self::is_compound(&values[0]) {
652 self.decision_over(
653 values[0].range(),
654 "value",
655 true,
656 |this, leaves, boolops| this.tree(&values[0], leaves, boolops).map(|_| ()),
657 );
658 }
659 }
660 Expr::UnaryOp(unary) if unary.op == UnaryOp::Not => {
661 let (operand, _) = Self::strip_not(expression);
662 let Expr::BoolOp(boolean) = operand else {
663 return;
664 };
665 if bounds(self, boolean.range)
666 .is_some_and(|own| self.decision_boolops.contains_key(&own))
667 {
668 return;
669 }
670 self.decision_over(unary.range, "value", true, |this, leaves, boolops| {
671 this.tree(expression, leaves, boolops).map(|_| ())
672 });
673 }
674 _ => {}
675 }
676 }
677
678 fn is_compound(expression: &Expr) -> bool {
679 matches!(Self::strip_not(expression).0, Expr::BoolOp(_))
680 }
681
682 fn branch<const N: usize>(
683 &mut self,
684 range: TextRange,
685 kind: &str,
686 alternatives: [(&str, &str); N],
687 ) -> Option<String> {
688 let id = stable_id(self.file, "branch", range, kind);
689 let (_, _, source) = self.location_source(range)?;
690 self.branch_with_id(id, range, kind, source, alternatives)
691 }
692
693 fn branch_with_id<const N: usize>(
694 &mut self,
695 id: String,
696 range: TextRange,
697 kind: &str,
698 source: String,
699 alternatives: [(&str, &str); N],
700 ) -> Option<String> {
701 if !self.branch_ids.insert(id.clone()) {
702 return None;
703 }
704 let (line, column, _) = self.location_source(range)?;
705 self.manifest.branches.push(BranchMeta {
706 id: id.clone(),
707 kind: kind.into(),
708 file: self.file.into(),
709 line,
710 column,
711 source,
712 alternatives: alternatives
713 .into_iter()
714 .map(|(suffix, label)| BranchAlternativeMeta {
715 id: format!("{id}:{suffix}"),
716 label: label.into(),
717 })
718 .collect(),
719 });
720 Some(id)
721 }
722
723 fn loop_branch(&mut self, range: TextRange, iter: &Expr, kind: &str) {
724 let Some(id) = self.branch(
725 range,
726 kind,
727 [("zero", "zero iterations"), ("entered", "entered")],
728 ) else {
729 return;
730 };
731 let Some(iter_span) = self.span(iter.range()) else {
732 return;
733 };
734 self.plan.loops.push(LoopPlan {
735 zero: format!("{id}:zero"),
736 entered: format!("{id}:entered"),
737 id,
738 iter: iter_span,
739 first: None,
740 });
741 }
742
743 fn try_statement(&mut self, statement: &ruff_python_ast::StmtTry) {
744 let Some(id) = self.branch(
745 statement.range,
746 if statement.is_star { "try-star" } else { "try" },
747 [("success", "try completed"), ("raised", "handler entered")],
748 ) else {
749 return;
750 };
751 let body_range = match (statement.body.first(), statement.body.last()) {
752 (Some(first), Some(last)) => TextRange::new(first.start(), last.end()),
753 _ => return,
754 };
755 let Some(body) = self.span(body_range) else {
756 return;
757 };
758 let block_span = |collector: &mut Self, block: &[Stmt]| -> Option<Option<PlanSpan>> {
759 match (block.first(), block.last()) {
760 (Some(first), Some(last)) => collector
761 .span(TextRange::new(first.start(), last.end()))
762 .map(Some),
763 _ => Some(None),
764 }
765 };
766 let Some(finalbody) = block_span(self, &statement.finalbody) else {
767 return;
768 };
769 let mut handlers = Vec::with_capacity(statement.handlers.len());
770 for (index, handler) in statement.handlers.iter().enumerate() {
771 let Some(handler_id) = self.branch(
772 handler.range(),
773 &format!("except-{index}"),
774 [("missed", "not selected"), ("selected", "selected")],
775 ) else {
776 return;
777 };
778 let ruff_python_ast::ExceptHandler::ExceptHandler(clause) = handler;
779 handlers.push(HandlerPlan {
780 missed: format!("{handler_id}:missed"),
781 selected: format!("{handler_id}:selected"),
782 id: handler_id,
783 bare: clause.type_.is_none(),
784 });
785 }
786 self.plan.tries.push(TryPlan {
787 success: format!("{id}:success"),
788 raised: format!("{id}:raised"),
789 id,
790 body,
791 finalbody,
792 handlers,
793 });
794 }
795
796 fn logical(&mut self, boolean: &ruff_python_ast::ExprBoolOp) {
797 let Some(boolop_span) = self.span(boolean.range) else {
798 return;
799 };
800 let bounds = match self.locations.range(boolean.range) {
801 Ok(bounds) => bounds,
802 Err(error) => {
803 self.fail::<()>(error);
804 return;
805 }
806 };
807 let decision = self.decision_boolops.get(&bounds).cloned();
808 let regrouped = self.regrouped.get(&bounds).cloned();
809 let op = match boolean.op {
810 BoolOp::And => "and",
811 BoolOp::Or => "or",
812 };
813 let last = boolean.values.len() - 1;
814 for (index, value) in boolean.values.iter().enumerate().skip(1) {
815 let Some(id) = self.branch(
816 value.range(),
817 &format!("logical-{op}-{index}"),
818 [
819 ("short-circuit", "short-circuited"),
820 ("evaluated", "right operand evaluated"),
821 ],
822 ) else {
823 continue;
824 };
825 let (decision_id, previous_leaves, operand_leaves, boolop, operand) =
829 match (&decision, ®rouped) {
830 (Some((decision_id, groups)), _) => (
831 Some(decision_id.clone()),
832 Some(groups[index - 1].clone()),
833 Some(groups[index].clone()),
834 boolop_span,
835 index,
836 ),
837 (None, Some((decision_id, groups, inner))) if index < last => (
838 Some(decision_id.clone()),
839 Some(groups[index - 1].clone()),
840 Some(groups[index].clone()),
841 *inner,
842 index,
843 ),
844 (None, Some(_)) => (None, None, None, boolop_span, 1),
845 (None, None) => (None, None, None, boolop_span, index),
846 };
847 self.plan.logical.push(LogicalPlan {
848 short_circuit: format!("{id}:short-circuit"),
849 evaluated: format!("{id}:evaluated"),
850 id,
851 boolop,
852 operand,
853 decision: decision_id,
854 previous_leaves,
855 operand_leaves,
856 });
857 }
858 }
859
860 fn match_statement(&mut self, statement: &ruff_python_ast::StmtMatch) {
861 let Some(span) = self.span(statement.range) else {
862 return;
863 };
864 let mut cases = Vec::with_capacity(statement.cases.len());
865 for (index, case) in statement.cases.iter().enumerate() {
866 let kind = format!("match-case-{index}");
867 let Some(id) = self.branch(
868 case.range,
869 &kind,
870 [("missed", "not selected"), ("selected", "selected")],
871 ) else {
872 return;
873 };
874 if let Some(guard) = &case.guard {
875 self.decision(guard, "match-guard");
876 }
877 let Some(case_span) = self.span(case.range) else {
878 return;
879 };
880 cases.push(MatchCasePlan {
881 missed: format!("{id}:missed"),
882 selected: format!("{id}:selected"),
883 id,
884 span: case_span,
885 irrefutable: case.guard.is_none() && case.pattern.is_irrefutable(),
886 });
887 }
888 let no_case = if statement
889 .cases
890 .iter()
891 .any(|case| case.guard.is_none() && case.pattern.is_irrefutable())
892 {
893 None
894 } else {
895 self.branch(
896 statement.subject.range(),
897 "match-no-case",
898 [
899 ("matched", "some case matched"),
900 ("unmatched", "no case matched"),
901 ],
902 )
903 .map(|id| MatchNoCasePlan {
904 matched: format!("{id}:matched"),
905 unmatched: format!("{id}:unmatched"),
906 id,
907 })
908 };
909 self.plan.matches.push(MatchPlan {
910 span,
911 cases,
912 no_case,
913 });
914 }
915
916 fn first_statement(&self, body: &[Stmt]) -> Option<String> {
919 body.iter()
920 .enumerate()
921 .find(|(index, statement)| !is_unobservable_statement(statement, *index == 0))
922 .map(|(_, statement)| stable_id(self.file, "statement", statement.range(), ""))
923 }
924
925 fn else_head(&self, clauses: &[ruff_python_ast::ElifElseClause]) -> Option<String> {
929 clauses
930 .first()
931 .filter(|clause| clause.test.is_none())
932 .and_then(|clause| self.first_statement(&clause.body))
933 }
934
935 fn branch_heads(&mut self, test: &Expr, when_true: Option<String>, when_false: Option<String>) {
938 let Ok(span) = self.locations.span(test.range()) else {
939 return;
940 };
941 if let Some(plan) = self
942 .plan
943 .decisions
944 .iter_mut()
945 .rev()
946 .find(|plan| plan.span == span)
947 {
948 plan.when_true = when_true;
949 plan.when_false = when_false;
950 }
951 }
952
953 fn visit_body_statements(&mut self, body: &'a [Stmt]) {
954 for (index, statement) in body.iter().enumerate() {
955 if is_unobservable_statement(statement, index == 0) {
956 continue;
959 }
960 self.visit_stmt(statement);
961 }
962 }
963}
964
965impl<'a> Visitor<'a> for PythonObligationCollector<'a> {
966 fn visit_body(&mut self, body: &'a [Stmt]) {
967 self.visit_body_statements(body);
968 }
969
970 fn visit_stmt(&mut self, statement: &'a Stmt) {
971 self.statement(statement);
972 match statement {
973 Stmt::FunctionDef(function) => {
974 self.function(function.range, &function.name);
975 let first = self.first_statement(&function.body);
976 if let Some(plan) = self.plan.functions.last_mut()
977 && plan.id == stable_id(self.file, "function", function.range, &function.name)
978 {
979 plan.first = first;
980 }
981 }
982 Stmt::If(statement) => {
983 let heads = (
984 self.first_statement(&statement.body),
985 self.else_head(&statement.elif_else_clauses),
986 );
987 self.decision(&statement.test, "if");
988 self.branch_heads(&statement.test, heads.0, heads.1);
989 for (index, clause) in statement.elif_else_clauses.iter().enumerate() {
990 if let Some(test) = &clause.test {
991 let heads = (
992 self.first_statement(&clause.body),
993 self.else_head(&statement.elif_else_clauses[index + 1..]),
994 );
995 self.decision(test, "elif");
996 self.branch_heads(test, heads.0, heads.1);
997 }
998 }
999 }
1000 Stmt::While(statement) => {
1001 let heads = (
1002 self.first_statement(&statement.body),
1003 self.first_statement(&statement.orelse),
1004 );
1005 self.decision(&statement.test, "while");
1006 self.branch_heads(&statement.test, heads.0, heads.1);
1007 }
1008 Stmt::For(statement) => {
1009 self.loop_branch(
1010 statement.range,
1011 &statement.iter,
1012 if statement.is_async {
1013 "async-for"
1014 } else {
1015 "for"
1016 },
1017 );
1018 let first = self.first_statement(&statement.body);
1019 if let (Some(plan), Ok(span)) = (
1020 self.plan.loops.last_mut(),
1021 self.locations.span(statement.iter.range()),
1022 ) && plan.iter == span
1023 {
1024 plan.first = first;
1025 }
1026 }
1027 Stmt::Match(statement) => self.match_statement(statement),
1028 Stmt::Try(statement) => self.try_statement(statement),
1029 Stmt::Assert(statement) => {
1030 self.decision(&statement.test, "assert");
1031 }
1032 _ => {}
1033 }
1034 walk_stmt(self, statement);
1035 }
1036
1037 fn visit_expr(&mut self, expression: &'a Expr) {
1038 match expression {
1039 Expr::Lambda(lambda) => self.function(lambda.range, "<lambda>"),
1040 Expr::If(expression) => {
1041 self.decision(&expression.test, "ternary");
1042 }
1043 Expr::BoolOp(boolean) => {
1044 self.value_decision(expression);
1045 self.logical(boolean);
1046 }
1047 Expr::UnaryOp(_) => self.value_decision(expression),
1048 _ => {}
1049 }
1050 walk_expr(self, expression);
1051 }
1052
1053 fn visit_comprehension(&mut self, comprehension: &'a Comprehension) {
1054 self.loop_branch(
1055 comprehension.range,
1056 &comprehension.iter,
1057 if comprehension.is_async {
1058 "async-comprehension"
1059 } else {
1060 "comprehension"
1061 },
1062 );
1063 for condition in &comprehension.ifs {
1064 self.decision(condition, "comprehension-if");
1065 }
1066 walk_comprehension(self, comprehension);
1067 }
1068}
1069
1070pub fn build_python_obligations(
1074 file: &str,
1075 source: &str,
1076) -> Result<PythonFileObligations, PythonInstrumenterError> {
1077 if source.len() > u32::MAX as usize {
1078 return Err(PythonInstrumenterError::SourceTooLarge);
1079 }
1080 let parsed =
1081 parse_module(source).map_err(|error| PythonInstrumenterError::Parse(error.to_string()))?;
1082 let mut collector = PythonObligationCollector::new(file, source);
1083 collector.visit_body(parsed.suite());
1084 if let Some(error) = collector.error {
1085 return Err(error);
1086 }
1087 collector.manifest.unmeasured.sort();
1088 collector.manifest.unmeasured.dedup();
1089 Ok(PythonFileObligations {
1090 manifest: collector.manifest,
1091 plan: collector.plan,
1092 })
1093}
1094
1095pub fn build_python_manifest(
1097 file: &str,
1098 source: &str,
1099) -> Result<CoverageManifest, PythonInstrumenterError> {
1100 build_python_obligations(file, source).map(|obligations| obligations.manifest)
1101}
1102
1103#[cfg(test)]
1104mod tests {
1105 use super::*;
1106
1107 const SOURCE: &str = r#"async def classify[T](items, flag=True):
1108 values = [item async for item in items if item.ready and flag]
1109 for value in values:
1110 if value.primary and (value.safe or flag):
1111 return value if flag else None
1112 elif value.fallback:
1113 break
1114 try:
1115 match values:
1116 case [first, *_] if first.ready:
1117 return first
1118 case []:
1119 return None
1120 except* ValueError:
1121 return None
1122 return (lambda value: value or flag)(None)
1123"#;
1124
1125 #[test]
1126 fn discovers_current_python_obligations_with_exact_ranges_and_stable_ids() {
1127 let first = build_python_obligations("src/app.py", SOURCE).unwrap();
1128 let second = build_python_obligations("src/app.py", SOURCE).unwrap();
1129 assert_eq!(first, second);
1130 let manifest = &first.manifest;
1131 assert!(manifest.points.iter().any(|point| {
1132 point.kind == PointKind::Function && point.label.as_deref() == Some("classify")
1133 }));
1134 assert!(manifest.points.iter().any(|point| {
1135 point.kind == PointKind::Function && point.label.as_deref() == Some("<lambda>")
1136 }));
1137 let if_decision = manifest
1138 .decisions
1139 .iter()
1140 .find(|decision| decision.kind == "if")
1141 .unwrap();
1142 assert_eq!(
1143 if_decision.conditions,
1144 ["value.primary", "value.safe", "flag"]
1145 );
1146 assert_eq!(if_decision.line, 4);
1147 assert_eq!(if_decision.column, 11);
1148 for kind in ["ternary", "match-guard", "comprehension-if", "elif"] {
1149 assert!(
1150 manifest
1151 .decisions
1152 .iter()
1153 .any(|decision| decision.kind == kind),
1154 "missing {kind}"
1155 );
1156 }
1157 assert!(
1158 manifest
1159 .branches
1160 .iter()
1161 .any(|branch| branch.kind == "async-comprehension")
1162 );
1163 assert!(
1164 manifest
1165 .branches
1166 .iter()
1167 .any(|branch| branch.kind == "try-star")
1168 );
1169 assert!(
1170 manifest
1171 .branches
1172 .iter()
1173 .any(|branch| branch.kind.starts_with("logical-and"))
1174 );
1175 assert!(
1176 manifest
1177 .decisions
1178 .iter()
1179 .all(|decision| decision.id.starts_with("py:decision:"))
1180 );
1181 assert!(manifest.unmeasured.is_empty());
1182 assert!(manifest.limitations.is_empty());
1183 }
1184
1185 #[test]
1186 fn a_compound_left_part_of_a_value_is_a_decision_whose_probes_keep_values() {
1187 let source = "def pick(a, b, c):\n first = a or b or c\n return (a and b) or c, not (a and b), a or c\n";
1188 let output = build_python_obligations("src/app.py", source).unwrap();
1189 let values = output
1190 .plan
1191 .decisions
1192 .iter()
1193 .filter(|decision| decision.kind == "value")
1194 .map(|decision| {
1195 (
1196 decision.span.start,
1197 decision.span.end,
1198 decision.conditions.len(),
1199 decision.value,
1200 )
1201 })
1202 .collect::<Vec<_>>();
1203 assert_eq!(
1206 values,
1207 [
1208 ([2, 12], [2, 18], 2, true),
1209 ([3, 12], [3, 19], 2, true),
1210 ([3, 27], [3, 40], 2, true),
1211 ]
1212 );
1213 let first = output
1216 .plan
1217 .logical
1218 .iter()
1219 .filter(|logical| logical.boolop.start == [2, 12])
1220 .map(|logical| {
1221 (
1222 logical.operand,
1223 logical.decision.is_some(),
1224 logical.boolop.end,
1225 )
1226 })
1227 .collect::<Vec<_>>();
1228 assert_eq!(first, [(1, true, [2, 18]), (1, false, [2, 23])]);
1229 }
1230
1231 #[test]
1232 fn plan_carries_spans_polarity_and_logical_leaves() {
1233 let plan = build_python_obligations("src/app.py", SOURCE).unwrap().plan;
1234 let if_plan = plan
1235 .decisions
1236 .iter()
1237 .find(|decision| decision.kind == "if")
1238 .unwrap();
1239 assert_eq!(if_plan.conditions.len(), 3);
1240 assert_eq!(if_plan.conditions[0].span.start, [4, 11]);
1241 assert_eq!(if_plan.conditions[0].span.end, [4, 24]);
1242 assert_eq!(if_plan.conditions[1].span.start, [4, 30]);
1243 let mut in_decision = plan
1247 .logical
1248 .iter()
1249 .filter(|logical| logical.decision.as_deref() == Some(if_plan.id.as_str()))
1250 .map(|logical| {
1251 (
1252 logical.previous_leaves.clone().unwrap(),
1253 logical.operand_leaves.clone().unwrap(),
1254 )
1255 })
1256 .collect::<Vec<_>>();
1257 in_decision.sort();
1258 assert_eq!(in_decision, [(vec![0], vec![1, 2]), (vec![1], vec![2])]);
1259 assert!(
1261 plan.logical
1262 .iter()
1263 .any(|logical| logical.decision.is_none())
1264 );
1265 assert!(
1266 plan.decisions
1267 .iter()
1268 .any(|decision| decision.kind == "comprehension-if")
1269 );
1270 assert!(
1271 plan.statements
1272 .iter()
1273 .any(|statement| statement.start == [1, 0])
1274 );
1275 assert!(
1276 plan.functions
1277 .iter()
1278 .any(|function| function.name == "classify" && function.span.start[0] == 1)
1279 );
1280 assert_eq!(plan.loops.len(), 2);
1281 assert_eq!(plan.matches.len(), 1);
1282 assert_eq!(plan.matches[0].cases.len(), 2);
1283 assert!(plan.matches[0].no_case.is_some());
1284 assert_eq!(plan.tries.len(), 1);
1285 let try_plan = &plan.tries[0];
1286 assert_eq!(try_plan.body.start, [9, 8]);
1287 assert_eq!(try_plan.handlers.len(), 1);
1288 assert!(!try_plan.handlers[0].bare);
1289 assert!(try_plan.finalbody.is_none());
1290 }
1291
1292 #[test]
1293 fn not_polarity_and_same_line_statements_are_modelled() {
1294 let source = "def f(a, b):\n if not (a and b): return 1\n x = 1; y = 2\n g = lambda: 1; h = lambda: 2\n return x + y\n";
1295 let obligations = build_python_obligations("m.py", source).unwrap();
1296 let decision = &obligations.plan.decisions[0];
1297 assert_eq!(decision.conditions.len(), 2);
1300 assert_eq!(decision.conditions[0].not, 0);
1301 assert_eq!(obligations.manifest.decisions[0].conditions, ["a", "b"]);
1302 let negated_leaf = build_python_obligations(
1303 "n.py",
1304 "def g(a):
1305 if not a:
1306 return 1
1307",
1308 )
1309 .unwrap();
1310 assert_eq!(negated_leaf.plan.decisions[0].conditions[0].not, 1);
1311 assert_eq!(negated_leaf.manifest.decisions[0].conditions, ["not a"]);
1312 let starts = obligations
1316 .plan
1317 .statements
1318 .iter()
1319 .map(|statement| statement.start)
1320 .collect::<Vec<_>>();
1321 for shared in [[2, 22], [3, 11], [4, 19]] {
1322 assert!(starts.contains(&shared), "{shared:?} in {starts:?}");
1323 }
1324 assert!(obligations.manifest.unmeasured.is_empty());
1325 assert!(obligations.manifest.limitations.is_empty());
1326 let lambdas = obligations
1327 .plan
1328 .functions
1329 .iter()
1330 .filter(|function| function.name == "<lambda>")
1331 .collect::<Vec<_>>();
1332 assert_eq!(lambdas.len(), 2);
1333 assert_ne!(lambdas[0].span, lambdas[1].span);
1334 }
1335
1336 #[test]
1337 fn try_statements_carry_bodies_handlers_and_finally_spans() {
1338 let source = "def g(x):\n try:\n y = int(x)\n except ValueError:\n y = -1\n except:\n y = -2\n else:\n y += 1\n finally:\n x = None\n return y\n";
1339 let plan = build_python_obligations("t.py", source).unwrap().plan;
1340 let try_plan = &plan.tries[0];
1341 assert_eq!(
1342 try_plan.body,
1343 PlanSpan {
1344 start: [3, 8],
1345 end: [3, 18]
1346 }
1347 );
1348 assert_eq!(
1349 try_plan.finalbody,
1350 Some(PlanSpan {
1351 start: [11, 8],
1352 end: [11, 16]
1353 })
1354 );
1355 assert_eq!(try_plan.handlers.len(), 2);
1356 assert!(!try_plan.handlers[0].bare);
1357 assert!(try_plan.handlers[1].bare);
1358 }
1359
1360 #[test]
1361 fn docstrings_and_scope_declarations_are_not_statements() {
1362 let source = "\"\"\"module doc\"\"\"\nX = 1\ndef f():\n \"\"\"doc\"\"\"\n global X\n pass\n";
1363 let obligations = build_python_obligations("m.py", source).unwrap();
1364 let statements = obligations
1365 .manifest
1366 .points
1367 .iter()
1368 .filter(|point| point.kind == PointKind::Statement)
1369 .map(|point| point.line)
1370 .collect::<Vec<_>>();
1371 assert_eq!(statements, [2, 3, 6]);
1372 }
1373
1374 #[test]
1375 fn each_structure_names_the_statement_that_heads_it() {
1376 let source = "def f(a):\n \"\"\"doc\"\"\"\n global G\n x = 1\n if a:\n y = 2\n elif a > 1:\n y = 3\n else:\n y = 4\n while a:\n a -= 1\n else:\n z = 5\n for item in a:\n w = item\n if x:\n pass\n";
1379 let obligations = build_python_obligations("m.py", source).unwrap();
1380 let statement_on = |line: usize| {
1381 let start = obligations
1382 .plan
1383 .statements
1384 .iter()
1385 .find(|statement| statement.start[0] == line)
1386 .unwrap_or_else(|| panic!("a statement on line {line}"));
1387 Some(start.id.clone())
1388 };
1389 assert_eq!(obligations.plan.functions[0].first, statement_on(4));
1390 let decision = |line: usize| {
1391 obligations
1392 .plan
1393 .decisions
1394 .iter()
1395 .find(|decision| decision.span.start[0] == line)
1396 .unwrap()
1397 };
1398 assert_eq!(
1400 (
1401 decision(5).when_true.clone(),
1402 decision(5).when_false.clone()
1403 ),
1404 (statement_on(6), None)
1405 );
1406 assert_eq!(
1407 (
1408 decision(7).when_true.clone(),
1409 decision(7).when_false.clone()
1410 ),
1411 (statement_on(8), statement_on(10))
1412 );
1413 assert_eq!(
1414 (
1415 decision(11).when_true.clone(),
1416 decision(11).when_false.clone()
1417 ),
1418 (statement_on(12), statement_on(14))
1419 );
1420 assert_eq!(
1421 (
1422 decision(17).when_true.clone(),
1423 decision(17).when_false.clone()
1424 ),
1425 (statement_on(18), None)
1426 );
1427 assert_eq!(obligations.plan.loops[0].first, statement_on(16));
1428 }
1429
1430 #[test]
1431 fn rejects_invalid_python_without_partial_obligations() {
1432 assert!(matches!(
1433 build_python_manifest("src/broken.py", "if :\n pass\n"),
1434 Err(PythonInstrumenterError::Parse(_))
1435 ));
1436 }
1437}