diff --git a/ast-pyo3/src/lib.rs b/ast-pyo3/src/lib.rs index ad26fa93..65a336b0 100644 --- a/ast-pyo3/src/lib.rs +++ b/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="", *, 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::(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) diff --git a/ast/asdl_rs.py b/ast/asdl_rs.py index 4b78639b..8a3bc3b5 100755 --- a/ast/asdl_rs.py +++ b/ast/asdl_rs.py @@ -1125,10 +1125,13 @@ 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 {{ @@ -1136,6 +1139,14 @@ def visitSum(self, sum, name, depth): }} }} }} + {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, ) diff --git a/ast/src/gen/located.rs b/ast/src/gen/located.rs index cc78b7e4..8f696756 100644 --- a/ast/src/gen/located.rs +++ b/ast/src/gen/located.rs @@ -5,42 +5,62 @@ pub type Mod = crate::generic::Mod; pub type ModModule = crate::generic::ModModule; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for ModModule { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for ModModule { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} pub type ModInteractive = crate::generic::ModInteractive; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for ModInteractive { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for ModInteractive { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} pub type ModExpression = crate::generic::ModExpression; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for ModExpression { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for ModExpression { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} pub type ModFunctionType = crate::generic::ModFunctionType; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for ModFunctionType { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for ModFunctionType { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} #[cfg(feature = "all-nodes-with-ranges")] impl Located for Mod { @@ -53,6 +73,17 @@ impl Located for Mod { } } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for Mod { + fn range_mut(&mut self) -> &mut SourceRange { + match self { + Self::Module(node) => node.range_mut(), + Self::Interactive(node) => node.range_mut(), + Self::Expression(node) => node.range_mut(), + Self::FunctionType(node) => node.range_mut(), + } + } +} pub type Stmt = crate::generic::Stmt; @@ -64,6 +95,12 @@ impl Located for StmtFunctionDef { } } +impl LocatedMut for StmtFunctionDef { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtAsyncFunctionDef = crate::generic::StmtAsyncFunctionDef; impl Located for StmtAsyncFunctionDef { @@ -72,6 +109,12 @@ impl Located for StmtAsyncFunctionDef { } } +impl LocatedMut for StmtAsyncFunctionDef { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtClassDef = crate::generic::StmtClassDef; impl Located for StmtClassDef { @@ -80,6 +123,12 @@ impl Located for StmtClassDef { } } +impl LocatedMut for StmtClassDef { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtReturn = crate::generic::StmtReturn; impl Located for StmtReturn { @@ -88,6 +137,12 @@ impl Located for StmtReturn { } } +impl LocatedMut for StmtReturn { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtDelete = crate::generic::StmtDelete; impl Located for StmtDelete { @@ -96,6 +151,12 @@ impl Located for StmtDelete { } } +impl LocatedMut for StmtDelete { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtAssign = crate::generic::StmtAssign; impl Located for StmtAssign { @@ -104,6 +165,12 @@ impl Located for StmtAssign { } } +impl LocatedMut for StmtAssign { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtAugAssign = crate::generic::StmtAugAssign; impl Located for StmtAugAssign { @@ -112,6 +179,12 @@ impl Located for StmtAugAssign { } } +impl LocatedMut for StmtAugAssign { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtAnnAssign = crate::generic::StmtAnnAssign; impl Located for StmtAnnAssign { @@ -120,6 +193,12 @@ impl Located for StmtAnnAssign { } } +impl LocatedMut for StmtAnnAssign { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtFor = crate::generic::StmtFor; impl Located for StmtFor { @@ -128,6 +207,12 @@ impl Located for StmtFor { } } +impl LocatedMut for StmtFor { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtAsyncFor = crate::generic::StmtAsyncFor; impl Located for StmtAsyncFor { @@ -136,6 +221,12 @@ impl Located for StmtAsyncFor { } } +impl LocatedMut for StmtAsyncFor { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtWhile = crate::generic::StmtWhile; impl Located for StmtWhile { @@ -144,6 +235,12 @@ impl Located for StmtWhile { } } +impl LocatedMut for StmtWhile { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtIf = crate::generic::StmtIf; impl Located for StmtIf { @@ -152,6 +249,12 @@ impl Located for StmtIf { } } +impl LocatedMut for StmtIf { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtWith = crate::generic::StmtWith; impl Located for StmtWith { @@ -160,6 +263,12 @@ impl Located for StmtWith { } } +impl LocatedMut for StmtWith { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtAsyncWith = crate::generic::StmtAsyncWith; impl Located for StmtAsyncWith { @@ -168,6 +277,12 @@ impl Located for StmtAsyncWith { } } +impl LocatedMut for StmtAsyncWith { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtMatch = crate::generic::StmtMatch; impl Located for StmtMatch { @@ -176,6 +291,12 @@ impl Located for StmtMatch { } } +impl LocatedMut for StmtMatch { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtRaise = crate::generic::StmtRaise; impl Located for StmtRaise { @@ -184,6 +305,12 @@ impl Located for StmtRaise { } } +impl LocatedMut for StmtRaise { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtTry = crate::generic::StmtTry; impl Located for StmtTry { @@ -192,6 +319,12 @@ impl Located for StmtTry { } } +impl LocatedMut for StmtTry { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtTryStar = crate::generic::StmtTryStar; impl Located for StmtTryStar { @@ -200,6 +333,12 @@ impl Located for StmtTryStar { } } +impl LocatedMut for StmtTryStar { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtAssert = crate::generic::StmtAssert; impl Located for StmtAssert { @@ -208,6 +347,12 @@ impl Located for StmtAssert { } } +impl LocatedMut for StmtAssert { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtImport = crate::generic::StmtImport; impl Located for StmtImport { @@ -216,6 +361,12 @@ impl Located for StmtImport { } } +impl LocatedMut for StmtImport { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtImportFrom = crate::generic::StmtImportFrom; impl Located for StmtImportFrom { @@ -224,6 +375,12 @@ impl Located for StmtImportFrom { } } +impl LocatedMut for StmtImportFrom { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtGlobal = crate::generic::StmtGlobal; impl Located for StmtGlobal { @@ -232,6 +389,12 @@ impl Located for StmtGlobal { } } +impl LocatedMut for StmtGlobal { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtNonlocal = crate::generic::StmtNonlocal; impl Located for StmtNonlocal { @@ -240,6 +403,12 @@ impl Located for StmtNonlocal { } } +impl LocatedMut for StmtNonlocal { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtExpr = crate::generic::StmtExpr; impl Located for StmtExpr { @@ -248,6 +417,12 @@ impl Located for StmtExpr { } } +impl LocatedMut for StmtExpr { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtPass = crate::generic::StmtPass; impl Located for StmtPass { @@ -256,6 +431,12 @@ impl Located for StmtPass { } } +impl LocatedMut for StmtPass { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtBreak = crate::generic::StmtBreak; impl Located for StmtBreak { @@ -264,6 +445,12 @@ impl Located for StmtBreak { } } +impl LocatedMut for StmtBreak { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type StmtContinue = crate::generic::StmtContinue; impl Located for StmtContinue { @@ -272,6 +459,12 @@ impl Located for StmtContinue { } } +impl LocatedMut for StmtContinue { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + impl Located for Stmt { fn range(&self) -> SourceRange { match self { @@ -306,6 +499,40 @@ impl Located for Stmt { } } +impl LocatedMut for Stmt { + fn range_mut(&mut self) -> &mut SourceRange { + match self { + Self::FunctionDef(node) => node.range_mut(), + Self::AsyncFunctionDef(node) => node.range_mut(), + Self::ClassDef(node) => node.range_mut(), + Self::Return(node) => node.range_mut(), + Self::Delete(node) => node.range_mut(), + Self::Assign(node) => node.range_mut(), + Self::AugAssign(node) => node.range_mut(), + Self::AnnAssign(node) => node.range_mut(), + Self::For(node) => node.range_mut(), + Self::AsyncFor(node) => node.range_mut(), + Self::While(node) => node.range_mut(), + Self::If(node) => node.range_mut(), + Self::With(node) => node.range_mut(), + Self::AsyncWith(node) => node.range_mut(), + Self::Match(node) => node.range_mut(), + Self::Raise(node) => node.range_mut(), + Self::Try(node) => node.range_mut(), + Self::TryStar(node) => node.range_mut(), + Self::Assert(node) => node.range_mut(), + Self::Import(node) => node.range_mut(), + Self::ImportFrom(node) => node.range_mut(), + Self::Global(node) => node.range_mut(), + Self::Nonlocal(node) => node.range_mut(), + Self::Expr(node) => node.range_mut(), + Self::Pass(node) => node.range_mut(), + Self::Break(node) => node.range_mut(), + Self::Continue(node) => node.range_mut(), + } + } +} + pub type Expr = crate::generic::Expr; pub type ExprBoolOp = crate::generic::ExprBoolOp; @@ -316,6 +543,12 @@ impl Located for ExprBoolOp { } } +impl LocatedMut for ExprBoolOp { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprNamedExpr = crate::generic::ExprNamedExpr; impl Located for ExprNamedExpr { @@ -324,6 +557,12 @@ impl Located for ExprNamedExpr { } } +impl LocatedMut for ExprNamedExpr { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprBinOp = crate::generic::ExprBinOp; impl Located for ExprBinOp { @@ -332,6 +571,12 @@ impl Located for ExprBinOp { } } +impl LocatedMut for ExprBinOp { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprUnaryOp = crate::generic::ExprUnaryOp; impl Located for ExprUnaryOp { @@ -340,6 +585,12 @@ impl Located for ExprUnaryOp { } } +impl LocatedMut for ExprUnaryOp { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprLambda = crate::generic::ExprLambda; impl Located for ExprLambda { @@ -348,6 +599,12 @@ impl Located for ExprLambda { } } +impl LocatedMut for ExprLambda { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprIfExp = crate::generic::ExprIfExp; impl Located for ExprIfExp { @@ -356,6 +613,12 @@ impl Located for ExprIfExp { } } +impl LocatedMut for ExprIfExp { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprDict = crate::generic::ExprDict; impl Located for ExprDict { @@ -364,6 +627,12 @@ impl Located for ExprDict { } } +impl LocatedMut for ExprDict { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprSet = crate::generic::ExprSet; impl Located for ExprSet { @@ -372,6 +641,12 @@ impl Located for ExprSet { } } +impl LocatedMut for ExprSet { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprListComp = crate::generic::ExprListComp; impl Located for ExprListComp { @@ -380,6 +655,12 @@ impl Located for ExprListComp { } } +impl LocatedMut for ExprListComp { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprSetComp = crate::generic::ExprSetComp; impl Located for ExprSetComp { @@ -388,6 +669,12 @@ impl Located for ExprSetComp { } } +impl LocatedMut for ExprSetComp { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprDictComp = crate::generic::ExprDictComp; impl Located for ExprDictComp { @@ -396,6 +683,12 @@ impl Located for ExprDictComp { } } +impl LocatedMut for ExprDictComp { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprGeneratorExp = crate::generic::ExprGeneratorExp; impl Located for ExprGeneratorExp { @@ -404,6 +697,12 @@ impl Located for ExprGeneratorExp { } } +impl LocatedMut for ExprGeneratorExp { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprAwait = crate::generic::ExprAwait; impl Located for ExprAwait { @@ -412,6 +711,12 @@ impl Located for ExprAwait { } } +impl LocatedMut for ExprAwait { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprYield = crate::generic::ExprYield; impl Located for ExprYield { @@ -420,6 +725,12 @@ impl Located for ExprYield { } } +impl LocatedMut for ExprYield { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprYieldFrom = crate::generic::ExprYieldFrom; impl Located for ExprYieldFrom { @@ -428,6 +739,12 @@ impl Located for ExprYieldFrom { } } +impl LocatedMut for ExprYieldFrom { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprCompare = crate::generic::ExprCompare; impl Located for ExprCompare { @@ -436,6 +753,12 @@ impl Located for ExprCompare { } } +impl LocatedMut for ExprCompare { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprCall = crate::generic::ExprCall; impl Located for ExprCall { @@ -444,6 +767,12 @@ impl Located for ExprCall { } } +impl LocatedMut for ExprCall { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprFormattedValue = crate::generic::ExprFormattedValue; impl Located for ExprFormattedValue { @@ -452,6 +781,12 @@ impl Located for ExprFormattedValue { } } +impl LocatedMut for ExprFormattedValue { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprJoinedStr = crate::generic::ExprJoinedStr; impl Located for ExprJoinedStr { @@ -460,6 +795,12 @@ impl Located for ExprJoinedStr { } } +impl LocatedMut for ExprJoinedStr { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprConstant = crate::generic::ExprConstant; impl Located for ExprConstant { @@ -468,6 +809,12 @@ impl Located for ExprConstant { } } +impl LocatedMut for ExprConstant { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprAttribute = crate::generic::ExprAttribute; impl Located for ExprAttribute { @@ -476,6 +823,12 @@ impl Located for ExprAttribute { } } +impl LocatedMut for ExprAttribute { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprSubscript = crate::generic::ExprSubscript; impl Located for ExprSubscript { @@ -484,6 +837,12 @@ impl Located for ExprSubscript { } } +impl LocatedMut for ExprSubscript { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprStarred = crate::generic::ExprStarred; impl Located for ExprStarred { @@ -492,6 +851,12 @@ impl Located for ExprStarred { } } +impl LocatedMut for ExprStarred { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprName = crate::generic::ExprName; impl Located for ExprName { @@ -500,6 +865,12 @@ impl Located for ExprName { } } +impl LocatedMut for ExprName { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprList = crate::generic::ExprList; impl Located for ExprList { @@ -508,6 +879,12 @@ impl Located for ExprList { } } +impl LocatedMut for ExprList { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprTuple = crate::generic::ExprTuple; impl Located for ExprTuple { @@ -516,6 +893,12 @@ impl Located for ExprTuple { } } +impl LocatedMut for ExprTuple { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type ExprSlice = crate::generic::ExprSlice; impl Located for ExprSlice { @@ -524,6 +907,12 @@ impl Located for ExprSlice { } } +impl LocatedMut for ExprSlice { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + impl Located for Expr { fn range(&self) -> SourceRange { match self { @@ -558,6 +947,40 @@ impl Located for Expr { } } +impl LocatedMut for Expr { + fn range_mut(&mut self) -> &mut SourceRange { + match self { + Self::BoolOp(node) => node.range_mut(), + Self::NamedExpr(node) => node.range_mut(), + Self::BinOp(node) => node.range_mut(), + Self::UnaryOp(node) => node.range_mut(), + Self::Lambda(node) => node.range_mut(), + Self::IfExp(node) => node.range_mut(), + Self::Dict(node) => node.range_mut(), + Self::Set(node) => node.range_mut(), + Self::ListComp(node) => node.range_mut(), + Self::SetComp(node) => node.range_mut(), + Self::DictComp(node) => node.range_mut(), + Self::GeneratorExp(node) => node.range_mut(), + Self::Await(node) => node.range_mut(), + Self::Yield(node) => node.range_mut(), + Self::YieldFrom(node) => node.range_mut(), + Self::Compare(node) => node.range_mut(), + Self::Call(node) => node.range_mut(), + Self::FormattedValue(node) => node.range_mut(), + Self::JoinedStr(node) => node.range_mut(), + Self::Constant(node) => node.range_mut(), + Self::Attribute(node) => node.range_mut(), + Self::Subscript(node) => node.range_mut(), + Self::Starred(node) => node.range_mut(), + Self::Name(node) => node.range_mut(), + Self::List(node) => node.range_mut(), + Self::Tuple(node) => node.range_mut(), + Self::Slice(node) => node.range_mut(), + } + } +} + pub type ExprContext = crate::generic::ExprContext; pub type ExprContextLoad = crate::generic::ExprContextLoad; @@ -635,12 +1058,17 @@ pub type CmpopNotIn = crate::generic::CmpopNotIn; pub type Comprehension = crate::generic::Comprehension; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for Comprehension { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for Comprehension { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} pub type Excepthandler = crate::generic::Excepthandler; @@ -652,6 +1080,12 @@ impl Located for ExcepthandlerExceptHandler { } } +impl LocatedMut for ExcepthandlerExceptHandler { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + impl Located for Excepthandler { fn range(&self) -> SourceRange { match self { @@ -660,15 +1094,28 @@ impl Located for Excepthandler { } } +impl LocatedMut for Excepthandler { + fn range_mut(&mut self) -> &mut SourceRange { + match self { + Self::ExceptHandler(node) => node.range_mut(), + } + } +} + pub type PythonArguments = crate::generic::PythonArguments; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for PythonArguments { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for PythonArguments { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} pub type Arg = crate::generic::Arg; @@ -678,6 +1125,12 @@ impl Located for Arg { } } +impl LocatedMut for Arg { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type Keyword = crate::generic::Keyword; impl Located for Keyword { @@ -686,6 +1139,12 @@ impl Located for Keyword { } } +impl LocatedMut for Keyword { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type Alias = crate::generic::Alias; impl Located for Alias { @@ -694,25 +1153,41 @@ impl Located for Alias { } } +impl LocatedMut for Alias { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type Withitem = crate::generic::Withitem; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for Withitem { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for Withitem { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} pub type MatchCase = crate::generic::MatchCase; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for MatchCase { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for MatchCase { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} pub type Pattern = crate::generic::Pattern; @@ -724,6 +1199,12 @@ impl Located for PatternMatchValue { } } +impl LocatedMut for PatternMatchValue { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type PatternMatchSingleton = crate::generic::PatternMatchSingleton; impl Located for PatternMatchSingleton { @@ -732,6 +1213,12 @@ impl Located for PatternMatchSingleton { } } +impl LocatedMut for PatternMatchSingleton { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type PatternMatchSequence = crate::generic::PatternMatchSequence; impl Located for PatternMatchSequence { @@ -740,6 +1227,12 @@ impl Located for PatternMatchSequence { } } +impl LocatedMut for PatternMatchSequence { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type PatternMatchMapping = crate::generic::PatternMatchMapping; impl Located for PatternMatchMapping { @@ -748,6 +1241,12 @@ impl Located for PatternMatchMapping { } } +impl LocatedMut for PatternMatchMapping { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type PatternMatchClass = crate::generic::PatternMatchClass; impl Located for PatternMatchClass { @@ -756,6 +1255,12 @@ impl Located for PatternMatchClass { } } +impl LocatedMut for PatternMatchClass { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type PatternMatchStar = crate::generic::PatternMatchStar; impl Located for PatternMatchStar { @@ -764,6 +1269,12 @@ impl Located for PatternMatchStar { } } +impl LocatedMut for PatternMatchStar { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type PatternMatchAs = crate::generic::PatternMatchAs; impl Located for PatternMatchAs { @@ -772,6 +1283,12 @@ impl Located for PatternMatchAs { } } +impl LocatedMut for PatternMatchAs { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + pub type PatternMatchOr = crate::generic::PatternMatchOr; impl Located for PatternMatchOr { @@ -780,6 +1297,12 @@ impl Located for PatternMatchOr { } } +impl LocatedMut for PatternMatchOr { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} + impl Located for Pattern { fn range(&self) -> SourceRange { match self { @@ -795,17 +1318,37 @@ impl Located for Pattern { } } +impl LocatedMut for Pattern { + fn range_mut(&mut self) -> &mut SourceRange { + match self { + Self::MatchValue(node) => node.range_mut(), + Self::MatchSingleton(node) => node.range_mut(), + Self::MatchSequence(node) => node.range_mut(), + Self::MatchMapping(node) => node.range_mut(), + Self::MatchClass(node) => node.range_mut(), + Self::MatchStar(node) => node.range_mut(), + Self::MatchAs(node) => node.range_mut(), + Self::MatchOr(node) => node.range_mut(), + } + } +} + pub type TypeIgnore = crate::generic::TypeIgnore; pub type TypeIgnoreTypeIgnore = crate::generic::TypeIgnoreTypeIgnore; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for TypeIgnoreTypeIgnore { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for TypeIgnoreTypeIgnore { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} #[cfg(feature = "all-nodes-with-ranges")] impl Located for TypeIgnore { @@ -815,23 +1358,41 @@ impl Located for TypeIgnore { } } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for TypeIgnore { + fn range_mut(&mut self) -> &mut SourceRange { + match self { + Self::TypeIgnore(node) => node.range_mut(), + } + } +} pub type Arguments = crate::generic::Arguments; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for Arguments { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for Arguments { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} pub type ArgWithDefault = crate::generic::ArgWithDefault; #[cfg(feature = "all-nodes-with-ranges")] - impl Located for ArgWithDefault { fn range(&self) -> SourceRange { self.range } } +#[cfg(feature = "all-nodes-with-ranges")] +impl LocatedMut for ArgWithDefault { + fn range_mut(&mut self) -> &mut SourceRange { + &mut self.range + } +} diff --git a/ast/src/located.rs b/ast/src/located.rs index 4987dc57..bf252617 100644 --- a/ast/src/located.rs +++ b/ast/src/located.rs @@ -13,6 +13,10 @@ pub trait Located { } } +pub trait LocatedMut: Located { + fn range_mut(&mut self) -> &mut SourceRange; +} + pub type Suite = Vec; pub use crate::builtin::*; diff --git a/ast/src/source_locator.rs b/ast/src/source_locator.rs index ef7e43eb..366c32bb 100644 --- a/ast/src/source_locator.rs +++ b/ast/src/source_locator.rs @@ -1,9 +1,11 @@ +use crate::Fold; use rustpython_parser_core::{ - source_code::{SourceLocation, SourceLocator, SourceRange}, + source_code::{LinearLocator, RandomLocator, SourceLocation, SourceRange}, text_size::TextRange, }; +use std::{convert::Infallible, unreachable}; -impl crate::fold::Fold for SourceLocator<'_> { +impl crate::fold::Fold for RandomLocator<'_> { type TargetU = SourceRange; type Error = std::convert::Infallible; type UserContext = SourceLocation; @@ -21,3 +23,266 @@ impl crate::fold::Fold for SourceLocator<'_> { Ok((start..end).into()) } } + +fn linear_locate_expr_joined_str( + locator: &mut LinearLocator<'_>, + node: crate::ExprJoinedStr, + location: SourceRange, +) -> Result, Infallible> { + let crate::ExprJoinedStr { range: _, values } = node; + + let mut located_values = Vec::with_capacity(values.len()); + for value in values.into_iter() { + let located = match value { + crate::Expr::Constant(constant) => { + let node = crate::ExprConstant { + range: location, + value: constant.value, + kind: constant.kind, + }; + crate::Expr::Constant(node) + } + crate::Expr::FormattedValue(formatted) => { + let node = crate::ExprFormattedValue { + range: location, + value: locator.fold(formatted.value)?, + conversion: formatted.conversion, + format_spec: formatted + .format_spec + .map(|spec| match *spec { + crate::Expr::JoinedStr(joined_str) => { + let node = + linear_locate_expr_joined_str(locator, joined_str, location)?; + Ok(crate::Expr::JoinedStr(node)) + } + expr => locator.fold(expr), + }) + .transpose()? + .map(Box::new), + }; + crate::Expr::FormattedValue(node) + } + _ => unreachable!("missing expr type for joined_str?"), + }; + located_values.push(located); + } + + Ok(crate::ExprJoinedStr { + range: location, + values: located_values, + }) +} + +impl crate::fold::Fold for LinearLocator<'_> { + type TargetU = SourceRange; + type Error = std::convert::Infallible; + type UserContext = SourceLocation; + + fn will_map_user(&mut self, user: &TextRange) -> Self::UserContext { + self.locate(user.start()) + } + + fn map_user( + &mut self, + user: TextRange, + start: Self::UserContext, + ) -> Result { + let end = self.locate(user.end()); + Ok((start..end).into()) + } + + fn fold_expr_dict( + &mut self, + node: crate::ExprDict, + ) -> Result, Self::Error> { + let crate::ExprDict { + range, + keys, + values, + } = node; + let context = self.will_map_user(&range); + assert_eq!(keys.len(), values.len()); + let mut located_keys = Vec::with_capacity(keys.len()); + let mut located_values = Vec::with_capacity(values.len()); + for (key, value) in keys.into_iter().zip(values.into_iter()) { + located_keys.push(self.fold(key)?); + located_values.push(self.fold(value)?); + } + let range = self.map_user(range, context)?; + Ok(crate::ExprDict { + range, + keys: located_keys, + values: located_values, + }) + } + + fn fold_expr_if_exp( + &mut self, + node: crate::ExprIfExp, + ) -> Result, Self::Error> { + let crate::ExprIfExp { + range, + test, + body, + orelse, + } = node; + let context = self.will_map_user(&range); + let body = self.fold(body)?; + let test = self.fold(test)?; + let orelse = self.fold(orelse)?; + let range = self.map_user(range, context)?; + Ok(crate::ExprIfExp { + range, + test, + body, + orelse, + }) + } + + fn fold_stmt_class_def( + &mut self, + node: crate::StmtClassDef, + ) -> Result, Self::Error> { + let crate::StmtClassDef { + name, + bases, + keywords, + body, + decorator_list, + range, + } = node; + let decorator_list = self.fold(decorator_list)?; + let context = self.will_map_user(&range); + + let name = self.fold(name)?; + let bases = self.fold(bases)?; + let keywords = self.fold(keywords)?; + let body = self.fold(body)?; + let range = self.map_user(range, context)?; + Ok(crate::StmtClassDef { + name, + bases, + keywords, + body, + decorator_list, + range, + }) + } + fn fold_stmt_function_def( + &mut self, + node: crate::StmtFunctionDef, + ) -> Result, Self::Error> { + let crate::StmtFunctionDef { + name, + args, + body, + decorator_list, + returns, + type_comment, + range, + } = node; + let decorator_list = self.fold(decorator_list)?; + let context = self.will_map_user(&range); + + let name = self.fold(name)?; + let args: Box> = self.fold(args)?; + let returns = self.fold(returns)?; + let body = self.fold(body)?; + let type_comment = self.fold(type_comment)?; + let range = self.map_user(range, context)?; + Ok(crate::StmtFunctionDef { + name, + args, + body, + decorator_list, + returns, + type_comment, + range, + }) + } + fn fold_stmt_async_function_def( + &mut self, + node: crate::StmtAsyncFunctionDef, + ) -> Result, Self::Error> { + let crate::StmtAsyncFunctionDef { + name, + args, + body, + decorator_list, + returns, + type_comment, + range, + } = node; + let decorator_list = self.fold(decorator_list)?; + let context = self.will_map_user(&range); + + let name = self.fold(name)?; + let args: Box> = self.fold(args)?; + let returns = self.fold(returns)?; + let body = self.fold(body)?; + let type_comment = self.fold(type_comment)?; + let range = self.map_user(range, context)?; + Ok(crate::StmtAsyncFunctionDef { + name, + args, + body, + decorator_list, + returns, + type_comment, + range, + }) + } + fn fold_expr_joined_str( + &mut self, + node: crate::ExprJoinedStr, + ) -> Result, Self::Error> { + let start = self.locate(node.range.start()); + let end = self.locate_only(node.range.end()); + let location = SourceRange::new(start, end); + linear_locate_expr_joined_str(self, node, location) + } + + fn fold_expr_call( + &mut self, + node: crate::ExprCall, + ) -> Result, Self::Error> { + let crate::ExprCall { + range, + func, + args, + keywords, + } = node; + let context = self.will_map_user(&range); + let func = self.fold(func)?; + let keywords = LinearLookaheadLocator(self).fold(keywords)?; + let args = self.fold(args)?; + let range = self.map_user(range, context)?; + Ok(crate::ExprCall { + range, + func, + args, + keywords, + }) + } +} + +struct LinearLookaheadLocator<'a, 'b>(&'b mut LinearLocator<'a>); + +impl crate::fold::Fold for LinearLookaheadLocator<'_, '_> { + type TargetU = SourceRange; + type Error = std::convert::Infallible; + type UserContext = SourceLocation; + + fn will_map_user(&mut self, user: &TextRange) -> Self::UserContext { + self.0.locate_only(user.start()) + } + + fn map_user( + &mut self, + user: TextRange, + start: Self::UserContext, + ) -> Result { + let end = self.0.locate_only(user.end()); + Ok((start..end).into()) + } +} diff --git a/core/src/source_code.rs b/core/src/source_code.rs index 8080da46..fbc24434 100644 --- a/core/src/source_code.rs +++ b/core/src/source_code.rs @@ -1,9 +1,10 @@ // re-export our public interface +use crate::text_size::{TextLen, TextSize}; pub use ruff_source_location::*; pub type LineNumber = OneIndexed; -#[derive(Debug, Copy, Clone)] +#[derive(Debug, Copy, Clone, Default)] pub struct SourceRange { pub start: SourceLocation, pub end: Option, @@ -31,12 +32,12 @@ impl From> for SourceRange { } /// Converts source code byte-offset to Python convention line and column numbers. -pub struct SourceLocator<'a> { +pub struct RandomLocator<'a> { pub source: &'a str, index: LineIndex, } -impl<'a> SourceLocator<'a> { +impl<'a> RandomLocator<'a> { #[inline] pub fn new(source: &'a str) -> Self { let index = LineIndex::from_source_text(source); @@ -65,6 +66,182 @@ impl<'a> SourceLocator<'a> { } } +/// Converts source code byte-offset to Python convention line and column numbers. +pub struct LinearLocator<'a> { + pub source: &'a str, + state: LinearLocatorState, + #[cfg(debug_assertions)] + index: LineIndex, +} + +struct LinearLocatorState { + line_start: TextSize, + line_end: Option, + line_number: OneIndexed, + cursor: TextSize, + is_ascii: bool, +} + +impl LinearLocatorState { + fn init(source: &str) -> Self { + let mut line_start = TextSize::default(); + if source.starts_with('\u{feff}') { + line_start += '\u{feff}'.text_len(); + } + let (line_end, is_ascii) = if let Some(nl) = source.find('\n') { + let is_ascii = source[..nl].is_ascii(); + (Some(TextSize::new(nl as u32 + 1)), is_ascii) + } else { + (None, source.is_ascii()) + }; + let line_number = OneIndexed::MIN; + Self { + line_start, + line_end, + line_number, + cursor: line_start, + is_ascii, + } + } + + fn new_line_start(&self, next_offset: TextSize) -> Option { + if let Some(new_line_start) = self.line_end { + if new_line_start <= next_offset { + return Some(new_line_start); + } + } + None + } +} + +impl<'a> LinearLocator<'a> { + // nl = newline + + #[inline] + pub fn new(source: &'a str) -> Self { + let state = LinearLocatorState::init(source); + Self { + source, + state, + #[cfg(debug_assertions)] + index: LineIndex::from_source_text(source), + } + } + + pub fn locate(&mut self, offset: crate::text_size::TextSize) -> SourceLocation { + debug_assert!( + self.state.cursor <= offset, + "{:?} -> {:?} {}", + self.state.cursor, + offset, + &self.source[offset.to_usize()..self.state.cursor.to_usize()] + ); + let (column, new_state) = self.locate_inner(offset); + if let Some(state) = new_state { + self.state = state; + } else { + self.state.cursor = offset; + } + SourceLocation { + row: self.state.line_number, + column, + } + } + + pub fn locate_only(&mut self, offset: crate::text_size::TextSize) -> SourceLocation { + let (column, new_state) = self.locate_inner(offset); + let state = new_state.as_ref().unwrap_or(&self.state); + SourceLocation { + row: state.line_number, + column, + } + } + + fn locate_inner( + &mut self, + offset: crate::text_size::TextSize, + ) -> (OneIndexed, Option) { + let (column, new_state) = if let Some(new_line_start) = self.state.new_line_start(offset) { + // not fit in current line + let focused = &self.source[new_line_start.to_usize()..offset.to_usize()]; + let (lines, line_start, column) = if let Some(last_newline) = focused.rfind('\n') { + let last_newline = new_line_start.to_usize() + last_newline; + let lines = self.source[self.state.cursor.to_usize()..last_newline] + .matches('\n') + .count() as u32 + + 1; // TODO: \r + let line_start = last_newline as u32 + 1; + let column = offset.to_u32() - line_start; + (lines, line_start, column) + } else { + let column = (offset - new_line_start).to_u32(); + (1, new_line_start.to_u32(), column) + }; + let line_number = self.state.line_number.saturating_add(lines); + let (line_end, is_ascii) = + if let Some(newline) = self.source[line_start as usize..].find('\n') { + let newline = line_start as usize + newline; + debug_assert_eq!(&self.source[newline..][..1], "\n"); + let is_ascii = self.source[line_start as usize..newline].is_ascii(); + (Some(TextSize::new(newline as u32 + 1)), is_ascii) + } else { + let is_ascii = self.source[line_start as usize..].is_ascii(); + (None, is_ascii) + }; + let line_start = TextSize::new(line_start); + let state = LinearLocatorState { + line_start, + line_end, + line_number, + cursor: offset, + is_ascii, + }; + (column, Some(state)) + } else { + let column = (offset - self.state.line_start).to_u32(); + (column, None) + }; + let state = new_state.as_ref().unwrap_or(&self.state); + let column = if state.is_ascii { + column + } else { + self.source[state.line_start.to_usize()..][..column as usize] + .chars() + .count() as u32 + }; + let column = OneIndexed::from_zero_indexed(column); + #[cfg(debug_assertions)] + { + let location = SourceLocation { + row: state.line_number, + column, + }; + let source_code = SourceCode::new(self.source, &self.index); + assert_eq!( + location, + source_code.source_location(offset), + "input: {} -> {} {}", + self.state.cursor.to_usize(), + offset.to_usize(), + &self.source[self.state.cursor.to_usize()..offset.to_usize()] + ); + } + (column, new_state) + } + + pub fn locate_error(&mut self, base: crate::error::BaseError) -> LocatedError + where + T: Into, + { + let location = self.locate(base.offset); + LocatedError { + error: base.error.into(), + location: Some(location), + source_path: base.source_path, + } + } +} + #[derive(Debug, PartialEq, Eq)] pub struct LocatedError { pub error: T, @@ -124,3 +301,36 @@ where Some(&self.error) } } + +#[test] +fn test_linear_locator() { + let source = r#" +123456789 +abcdefghi + +유니코드 + "# + .strip_prefix(char::is_whitespace) + .unwrap(); + let mut locator = LinearLocator::new(source); + let mut random_locator = RandomLocator::new(source); + + let mut test = |(row, col), offset| { + let input = TextSize::from(offset); + let expected: SourceLocation = SourceLocation { + row: OneIndexed::new(row).unwrap(), + column: OneIndexed::new(col).unwrap(), + }; + let actual = locator.locate(input); + let actual2 = random_locator.locate(input); + assert_eq!(expected, actual); + assert_eq!(expected, actual2); + }; + + test((1, 1), 0); + test((1, 6), 5); + test((1, 9), 8); + test((2, 1), 10); + test((4, 1), 21); + test((4, 3), 27); +}