rustpython-ruff_python_ast 0.15.8

Unofficial fork for RustPython
Documentation 
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

//! Utilities for manually traversing a Python AST.
use crate::{self as ast, AnyNodeRef, ExceptHandler, Stmt};

/// Given a [`Stmt`] and its parent, return the [`ast::Suite`] that contains the [`Stmt`].
pub fn suite<'a>(
    stmt: impl Into<AnyNodeRef<'a>>,
    parent: impl Into<AnyNodeRef<'a>>,
) -> Option<EnclosingSuite<'a>> {
    // TODO: refactor this to work without a parent, ie when `stmt` is at the top level
    let stmt = stmt.into();
    match parent.into() {
        AnyNodeRef::ModModule(ast::ModModule { body, .. }) => EnclosingSuite::new(body, stmt),
        AnyNodeRef::StmtFunctionDef(ast::StmtFunctionDef { body, .. }) => {
            EnclosingSuite::new(body, stmt)
        }
        AnyNodeRef::StmtClassDef(ast::StmtClassDef { body, .. }) => EnclosingSuite::new(body, stmt),
        AnyNodeRef::StmtFor(ast::StmtFor { body, orelse, .. }) => [body, orelse]
            .iter()
            .find_map(|suite| EnclosingSuite::new(suite, stmt)),
        AnyNodeRef::StmtWhile(ast::StmtWhile { body, orelse, .. }) => [body, orelse]
            .iter()
            .find_map(|suite| EnclosingSuite::new(suite, stmt)),
        AnyNodeRef::StmtIf(ast::StmtIf {
            body,
            elif_else_clauses,
            ..
        }) => [body]
            .into_iter()
            .chain(elif_else_clauses.iter().map(|clause| &clause.body))
            .find_map(|suite| EnclosingSuite::new(suite, stmt)),
        AnyNodeRef::StmtWith(ast::StmtWith { body, .. }) => EnclosingSuite::new(body, stmt),
        AnyNodeRef::StmtMatch(ast::StmtMatch { cases, .. }) => cases
            .iter()
            .map(|case| &case.body)
            .find_map(|body| EnclosingSuite::new(body, stmt)),
        AnyNodeRef::StmtTry(ast::StmtTry {
            body,
            handlers,
            orelse,
            finalbody,
            ..
        }) => [body, orelse, finalbody]
            .into_iter()
            .chain(
                handlers
                    .iter()
                    .filter_map(ExceptHandler::as_except_handler)
                    .map(|handler| &handler.body),
            )
            .find_map(|suite| EnclosingSuite::new(suite, stmt)),
        _ => None,
    }
}

pub struct EnclosingSuite<'a> {
    suite: &'a [Stmt],
    position: usize,
}

impl<'a> EnclosingSuite<'a> {
    pub fn new(suite: &'a [Stmt], stmt: AnyNodeRef<'a>) -> Option<Self> {
        let position = suite
            .iter()
            .position(|sibling| AnyNodeRef::ptr_eq(sibling.into(), stmt))?;

        Some(EnclosingSuite { suite, position })
    }

    pub fn next_sibling(&self) -> Option<&'a Stmt> {
        self.suite.get(self.position + 1)
    }

    pub fn next_siblings(&self) -> &'a [Stmt] {
        self.suite.get(self.position + 1..).unwrap_or_default()
    }

    pub fn previous_sibling(&self) -> Option<&'a Stmt> {
        self.suite.get(self.position.checked_sub(1)?)
    }
}

impl std::ops::Deref for EnclosingSuite<'_> {
    type Target = [Stmt];

    fn deref(&self) -> &Self::Target {
        self.suite
    }
}