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::{Bound, PyAny, prelude::PyAnyMethods};
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// Note: Direct PyAny implementation removed in favor of Bound<PyAny> implementation below
42
43/// Generic helper function for extracting position attributes from PyAny objects.
44fn extract_position_attr(obj: &Bound<PyAny>, attr_name: &str) -> Option<usize> {
45    obj.getattr(attr_name)
46        .ok()
47        .and_then(|attr| attr.extract().ok())
48}
49
50impl<'py> Node for &Bound<'py, PyAny> {
51    /// A method for getting the starting line number of the node. This may not exist for all node types.
52    fn lineno(&self) -> Option<usize> {
53        extract_position_attr(self, "lineno")
54    }
55
56    /// A method for getting the starting column of the node. This may not exist for all node types.
57    fn col_offset(&self) -> Option<usize> {
58        extract_position_attr(self, "col_offset")
59    }
60
61    /// A method for getting the ending line number of the node. This may not exist for all node types.
62    fn end_lineno(&self) -> Option<usize> {
63        extract_position_attr(self, "end_lineno")
64    }
65
66    /// A method for getting the ending column of the node. This may not exist for all node types.
67    fn end_col_offset(&self) -> Option<usize> {
68        extract_position_attr(self, "end_col_offset")
69    }
70}