Linear Locator (#46)

Author: youknowone

ast-pyo3/src/lib.rs

@@ -4,7 +4,7 @@ pub mod wrapper;
 
 pub use py_ast::{init, PyNode, ToPyAst};
 use pyo3::prelude::*;
-use rustpython_parser::ast::{source_code::SourceLocator, Fold};
+use rustpython_parser::ast::{source_code::LinearLocator, Fold};
 
 #[pyfunction]
 #[pyo3(signature = (source, filename="<unknown>", *, type_comments=false, locate=true))]
@@ -21,7 +21,7 @@ pub fn parse<'py>(
     let parsed = rustpython_parser::parse(source, rustpython_parser::Mode::Module, filename)
         .map_err(|e| PyErr::new::<pyo3::exceptions::PySyntaxError, _>(e.to_string()))?;
     if locate {
-        let parsed = SourceLocator::new(source).fold(parsed).unwrap();
+        let parsed = LinearLocator::new(source).fold(parsed).unwrap();
         parsed.module().unwrap().to_py_ast(py)
     } else {
         parsed.module().unwrap().to_py_ast(py)

ast/asdl_rs.py

@@ -1125,17 +1125,28 @@ def visitSum(self, sum, name, depth):
             self.emit_located_impl(variant_info)
 
         if not info.no_cfg(self.type_info):
-            self.emit('#[cfg(feature = "all-nodes-with-ranges")]', 0)
+            cfg = '#[cfg(feature = "all-nodes-with-ranges")]'
+        else:
+            cfg = ''
 
         self.emit(
             f"""
+            {cfg}
             impl Located for {info.full_type_name} {{
                 fn range(&self) -> SourceRange {{
                     match self {{
                         {sum_match_arms}
                     }}
                 }}
             }}
+            {cfg}
+            impl LocatedMut for {info.full_type_name} {{
+                fn range_mut(&mut self) -> &mut SourceRange {{
+                    match self {{
+                        {sum_match_arms.replace('range()', 'range_mut()')}
+                    }}
+                }}
+            }}
             """.lstrip(),
             0,
         )
@@ -1157,15 +1168,24 @@ def emit_type_alias(self, info):
 
     def emit_located_impl(self, info):
         if not info.no_cfg(self.type_info):
-            self.emit('#[cfg(feature = "all-nodes-with-ranges")]', 0)
+            cfg = '#[cfg(feature = "all-nodes-with-ranges")]'
+        else:
+            cfg = ''
 
         self.emit(
             f"""
+            {cfg}
             impl Located for {info.full_type_name} {{
                 fn range(&self) -> SourceRange {{
                     self.range
                 }}
             }}
+            {cfg}
+            impl LocatedMut for {info.full_type_name} {{
+                fn range_mut(&mut self) -> &mut SourceRange {{
+                    &mut self.range
+                }}
+            }}
             """,
             0,
         )