Skip to main content

scarf_python/
node.rs

1// =======================================================================
2// node.rs
3// =======================================================================
4//! A wrapper around [`scarf_syntax::Node`]
5
6use pyo3::{exceptions::PyIOError, prelude::*};
7use scarf_parser::PreprocessorCache;
8use scarf_syntax::*;
9use std::io::{Read, Seek};
10use std::{fs::File, ops::Range};
11
12/// A wrapper around [`std::ops::Range<usize>`]
13#[pyclass(eq, from_py_object, module = "scarf_python")]
14#[derive(Clone, PartialEq, Eq)]
15pub struct Bytes {
16    /// The start of the byte-span (inclusive)
17    #[pyo3(get, set)]
18    pub start: usize,
19    /// The end of the byte-span (exclusive)
20    #[pyo3(get, set)]
21    pub end: usize,
22}
23
24impl<T> From<Range<T>> for Bytes
25where
26    T: Into<usize>,
27{
28    fn from(value: Range<T>) -> Self {
29        Self {
30            start: value.start.into(),
31            end: value.end.into(),
32        }
33    }
34}
35
36impl<T> From<Bytes> for Range<T>
37where
38    T: From<usize>,
39{
40    fn from(value: Bytes) -> Self {
41        Self {
42            start: value.start.into(),
43            end: value.end.into(),
44        }
45    }
46}
47
48/// A wrapper around [`scarf_syntax::Span`], providing a location in the
49/// source code where a [`Node`] was found
50#[pyclass(eq, from_py_object, module = "scarf_python")]
51#[derive(Clone, PartialEq, Eq)]
52pub struct Span {
53    /// The source file
54    #[pyo3(get, set)]
55    pub file: String,
56    /// The byte-span within the source file
57    #[pyo3(get, set)]
58    pub bytes: Bytes,
59    /// The [`Span`] of the original token, if this is a text macro
60    pub expanded_from: Option<Box<Span>>,
61    /// The [`Span`] of the `` `include `` directive that produced this
62    /// one, if any
63    pub included_from: Option<Box<Span>>,
64}
65
66impl<'a> From<scarf_syntax::Span<'a>> for Span {
67    fn from(value: scarf_syntax::Span<'a>) -> Self {
68        Self {
69            file: value.file.to_string(),
70            bytes: value.bytes.into(),
71            expanded_from: match value.expanded_from {
72                Some(expanded_from_ref) => {
73                    Some(Box::new(expanded_from_ref.clone().into()))
74                }
75                None => None,
76            },
77            included_from: match value.included_from {
78                Some(included_from_ref) => {
79                    Some(Box::new(included_from_ref.clone().into()))
80                }
81                None => None,
82            },
83        }
84    }
85}
86
87impl<'a> Span {
88    /// Turn a [`Span`] into a [`scarf_syntax::Span`]
89    pub(crate) fn to_span(
90        &'a self,
91        cache: &'a PreprocessorCache<'a>,
92    ) -> scarf_syntax::Span<'a> {
93        scarf_syntax::Span {
94            file: &self.file,
95            bytes: self.bytes.clone().into(),
96            expanded_from: match &self.expanded_from {
97                Some(expanded_from_box) => {
98                    Some(expanded_from_box.as_ref().to_span_ref(cache))
99                }
100                None => None,
101            },
102            included_from: match &self.included_from {
103                Some(included_from_box) => {
104                    Some(included_from_box.as_ref().to_span_ref(cache))
105                }
106                None => None,
107            },
108        }
109    }
110}
111
112#[pymethods]
113impl Span {
114    // Get the corresponding text from a [`Span`]
115    #[getter]
116    pub fn text(&self) -> PyResult<String> {
117        let mut file = File::open(self.file.as_str())
118            .map_err(|err| PyIOError::new_err(err.to_string()))?;
119        file.seek(std::io::SeekFrom::Start(self.bytes.start as u64))
120            .map_err(|err| PyIOError::new_err(err.to_string()))?;
121        let mut byte_buf = vec![0_u8; self.bytes.end - self.bytes.start];
122        file.read_exact(&mut byte_buf)
123            .map_err(|err| PyIOError::new_err(err.to_string()))?;
124        String::from_utf8(byte_buf)
125            .map_err(|err| PyIOError::new_err(err.to_string()))
126    }
127}
128
129impl<'a> Span {
130    /// Turn a [`Span`] into a [`&scarf_syntax::Span`]
131    pub(crate) fn to_span_ref(
132        &'a self,
133        cache: &'a PreprocessorCache<'a>,
134    ) -> &'a scarf_syntax::Span<'a> {
135        let new_span = self.to_span(cache);
136        cache.retain_span(new_span)
137    }
138}
139
140/// A wrapper around [`scarf_syntax::Node`], providing a single CST node
141#[pyclass(eq, from_py_object, module = "scarf_python")]
142#[derive(Clone, PartialEq, Eq)]
143pub struct Node {
144    /// The name of the [`Node`] (see [`scarf_syntax::Node::name`])
145    #[pyo3(get, set)]
146    pub name: String,
147    /// The span of the [`Node`] (see [`scarf_syntax::Node::span`])
148    #[pyo3(get, set)]
149    pub span: Span,
150    /// All direct children of the [`Node`] in the CST (see [`scarf_syntax::Node::children`])
151    #[pyo3(get, set)]
152    pub children: Vec<Node>,
153}
154
155impl<'a, 'b> From<scarf_syntax::Node<'a, 'b>> for Node {
156    fn from(value: scarf_syntax::Node<'a, 'b>) -> Self {
157        Self {
158            name: value.name().to_string(),
159            span: value.span().into(),
160            children: value
161                .children()
162                .into_iter()
163                .map(|rust_node: scarf_syntax::Node<'_, '_>| rust_node.into())
164                .collect(),
165        }
166    }
167}
168
169#[pymethods]
170impl Node {
171    fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<NodeIterator>> {
172        Py::new(slf.py(), NodeIterator::new(slf.clone()))
173    }
174    // Get the corresponding text from a [`Node`]
175    #[getter]
176    pub fn text(&self) -> PyResult<String> {
177        self.span.text()
178    }
179}
180
181/// An iterator over [`Node`]s
182#[pyclass(eq, from_py_object, module = "scarf_python")]
183#[derive(Clone, PartialEq, Eq)]
184pub struct NodeIterator {
185    nodes: Vec<Node>,
186}
187
188impl NodeIterator {
189    pub fn new(root: Node) -> Self {
190        Self { nodes: vec![root] }
191    }
192}
193
194#[pymethods]
195impl NodeIterator {
196    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
197        slf
198    }
199    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<Node> {
200        match slf.nodes.pop() {
201            Some(node) => {
202                slf.nodes.append(
203                    &mut node
204                        .children
205                        .clone()
206                        .into_iter()
207                        .rev()
208                        .collect::<Vec<_>>(),
209                );
210                Some(node)
211            }
212            None => None,
213        }
214    }
215}