omnicoach 0.3.0

Lightweight, high-signal static analysis and command linter for Python and Rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

//! Header and decorator extents of Python statements.

use crate::code_lint::ast::tree::{ParsedFile, span_from_ruff_range};
use crate::diagnostic::SourceSpan;
use ruff_python_ast::Stmt;
use ruff_python_ast::token::TokenKind;
use ruff_python_ast::visitor::source_order::{SourceOrderVisitor, walk_stmt};
use ruff_text_size::Ranged as _;

/// Resolves the byte span of the header of the innermost Python statement enclosing `span`.
pub fn enclosing_statement_header_span(file: &ParsedFile, span: SourceSpan) -> Option<SourceSpan> {
    struct StatementFinder<'a> {
        target_span: SourceSpan,
        found: Option<&'a Stmt>,
    }

    impl<'a> SourceOrderVisitor<'a> for StatementFinder<'a> {
        fn visit_stmt(&mut self, statement: &'a Stmt) {
            if !span_from_ruff_range(statement.range()).contains(self.target_span) {
                return;
            }
            self.found = Some(statement);
            walk_stmt(self, statement);
        }
    }

    let parsed = file.py_module()?;
    let mut finder = StatementFinder {
        target_span: span,
        found: None,
    };
    finder.visit_body(&parsed.syntax().body);
    let statement = finder.found?;

    let statement_range = statement.range();
    let start_offset = statement_range.start().to_usize();
    let end_offset = python_statement_body_start(statement).map_or_else(
        || statement_range.end().to_usize(),
        |body_start| {
            parsed
                .tokens()
                .iter()
                .rfind(|token| {
                    token.range().start() >= statement_range.start()
                        && token.range().end() <= body_start
                        && !matches!(
                            token.kind(),
                            TokenKind::Comment
                                | TokenKind::Newline
                                | TokenKind::NonLogicalNewline
                                | TokenKind::Indent
                                | TokenKind::Dedent
                        )
                })
                .map_or(start_offset, |token| token.range().end().to_usize())
        },
    );
    Some(SourceSpan::new(start_offset, end_offset))
}

/// Returns the byte offset where the body of a Python compound `statement` begins, or `None`
/// for simple statements without a suite.
fn python_statement_body_start(statement: &Stmt) -> Option<ruff_text_size::TextSize> {
    match statement {
        Stmt::FunctionDef(node) => node.body.first().map(|first| first.range().start()),
        Stmt::ClassDef(node) => node.body.first().map(|first| first.range().start()),
        Stmt::For(node) => node.body.first().map(|first| first.range().start()),
        Stmt::While(node) => node.body.first().map(|first| first.range().start()),
        Stmt::If(node) => node.body.first().map(|first| first.range().start()),
        Stmt::With(node) => node.body.first().map(|first| first.range().start()),
        Stmt::Try(node) => node.body.first().map(|first| first.range().start()),
        Stmt::Match(node) => node.cases.first().map(|first| first.range().start()),
        _ => None,
    }
}

/// Records the byte span of every Python decorator in `file`.
pub fn record_decorator_and_attribute_spans(
    file: &ParsedFile,
    record: &mut impl FnMut(SourceSpan),
) {
    struct DecoratorVisitor<Record> {
        record: Record,
    }

    impl<'a, Record: FnMut(SourceSpan)> SourceOrderVisitor<'a> for DecoratorVisitor<Record> {
        fn visit_stmt(&mut self, statement: &'a Stmt) {
            let decorators = match statement {
                Stmt::FunctionDef(function) => function.decorator_list.as_slice(),
                Stmt::ClassDef(class) => class.decorator_list.as_slice(),
                _ => &[],
            };
            for decorator in decorators {
                (self.record)(span_from_ruff_range(decorator.range()));
            }
            walk_stmt(self, statement);
        }
    }

    if let Some(parsed) = file.py_module() {
        DecoratorVisitor { record }.visit_body(&parsed.syntax().body);
    }
}