python_ast/ast/
node.rs

1//! A module for AST elements that represent a position in a source file. Implementing the Node trait allows
2//! an ergonomic means of extracting line and column information from an item.
3
4use pyo3::{PyAny, PyResult};
5
6/// A trait for AST elements that represent a position in a source file. Implementing this trait allows
7/// an ergonomic means of extracting line and column information from an item.
8pub trait Node {
9    /// A method for getting the starting line number of the node. This may not exist for all node types.
10    fn lineno(&self) -> Option<usize> {
11        None
12    }
13
14    /// A method for getting the starting column of the node. This may not exist for all node types.
15    fn col_offset(&self) -> Option<usize> {
16        None
17    }
18
19    /// A method for getting the ending line number of the node. This may not exist for all node types.
20    fn end_lineno(&self) -> Option<usize> {
21        None
22    }
23
24    /// A method for getting the ending column of the node. This may not exist for all node types.
25    fn end_col_offset(&self) -> Option<usize> {
26        None
27    }
28
29    /// Generate an error message for the current code, adding line and column number.
30    fn error_message(&self, mod_name: impl AsRef<str>, message: impl AsRef<str>) -> String {
31        format!(
32            "{} {}:{:?}:{:?}",
33            message.as_ref(),
34            mod_name.as_ref(),
35            self.lineno(),
36            self.col_offset()
37        )
38    }
39}
40
41// These will only work on objects of Python's ast library's nodes, but you can try them on anything.
42impl Node for PyAny {
43    /// A method for getting the starting line number of the node. This may not exist for all node types.
44    fn lineno(&self) -> Option<usize> {
45        let lineno = self.getattr("lineno");
46        if let Ok(ln_any) = lineno {
47            let ln: PyResult<usize> = ln_any.extract();
48            if let Ok(l) = ln {
49                Some(l)
50            } else {
51                None
52            }
53        } else {
54            None
55        }
56    }
57
58    /// A method for getting the starting column of the node. This may not exist for all node types.
59    fn col_offset(&self) -> Option<usize> {
60        let col_offset = self.getattr("col_offset");
61        if let Ok(offset_any) = col_offset {
62            let ln: PyResult<usize> = offset_any.extract();
63            if let Ok(l) = ln {
64                Some(l)
65            } else {
66                None
67            }
68        } else {
69            None
70        }
71    }
72
73    /// A method for getting the ending line number of the node. This may not exist for all node types.
74    fn end_lineno(&self) -> Option<usize> {
75        let lineno = self.getattr("end_lineno");
76        if let Ok(ln_any) = lineno {
77            let ln: PyResult<usize> = ln_any.extract();
78            if let Ok(l) = ln {
79                Some(l)
80            } else {
81                None
82            }
83        } else {
84            None
85        }
86    }
87
88    /// A method for getting the ending column of the node. This may not exist for all node types.
89    fn end_col_offset(&self) -> Option<usize> {
90        let col_offset = self.getattr("end_col_offset");
91        if let Ok(offset_any) = col_offset {
92            let ln: PyResult<usize> = offset_any.extract();
93            if let Ok(l) = ln {
94                Some(l)
95            } else {
96                None
97            }
98        } else {
99            None
100        }
101    }
102}