1use crate::Fold;
2use rustpython_parser_core::{
3 source_code::{LinearLocator, RandomLocator, SourceLocation, SourceRange},
4 text_size::TextRange,
5};
6use std::{convert::Infallible, unreachable};
7
8impl crate::fold::Fold<TextRange> for RandomLocator<'_> {
9 type TargetU = SourceRange;
10 type Error = std::convert::Infallible;
11 type UserContext = SourceLocation;
12
13 fn will_map_user(&mut self, user: &TextRange) -> Self::UserContext {
14 self.locate(user.start())
15 }
16
17 fn map_user(
18 &mut self,
19 user: TextRange,
20 start: Self::UserContext,
21 ) -> Result<Self::TargetU, Self::Error> {
22 let end = self.locate(user.end());
23 Ok((start..end).into())
24 }
25}
26
27fn linear_locate_expr_joined_str(
28 locator: &mut LinearLocator<'_>,
29 node: crate::ExprJoinedStr<TextRange>,
30 location: SourceRange,
31) -> Result<crate::ExprJoinedStr<SourceRange>, Infallible> {
32 let crate::ExprJoinedStr { range: _, values } = node;
33
34 let mut located_values = Vec::with_capacity(values.len());
35 for value in values.into_iter() {
36 let located = match value {
37 crate::Expr::Constant(constant) => {
38 let node = crate::ExprConstant {
39 range: location,
40 value: constant.value,
41 kind: constant.kind,
42 };
43 crate::Expr::Constant(node)
44 }
45 crate::Expr::FormattedValue(formatted) => {
46 let node = crate::ExprFormattedValue {
47 range: location,
48 value: locator.fold(formatted.value)?,
49 conversion: formatted.conversion,
50 format_spec: formatted
51 .format_spec
52 .map(|spec| match *spec {
53 crate::Expr::JoinedStr(joined_str) => {
54 let node =
55 linear_locate_expr_joined_str(locator, joined_str, location)?;
56 Ok(crate::Expr::JoinedStr(node))
57 }
58 expr => locator.fold(expr),
59 })
60 .transpose()?
61 .map(Box::new),
62 };
63 crate::Expr::FormattedValue(node)
64 }
65 _ => unreachable!("missing expr type for joined_str?"),
66 };
67 located_values.push(located);
68 }
69
70 Ok(crate::ExprJoinedStr {
71 range: location,
72 values: located_values,
73 })
74}
75
76impl crate::fold::Fold<TextRange> for LinearLocator<'_> {
77 type TargetU = SourceRange;
78 type Error = std::convert::Infallible;
79 type UserContext = SourceLocation;
80
81 fn will_map_user(&mut self, user: &TextRange) -> Self::UserContext {
82 self.locate(user.start())
83 }
84
85 fn map_user(
86 &mut self,
87 user: TextRange,
88 start: Self::UserContext,
89 ) -> Result<Self::TargetU, Self::Error> {
90 let end = self.locate(user.end());
91 Ok((start..end).into())
92 }
93
94 fn fold_expr_dict(
95 &mut self,
96 node: crate::ExprDict<TextRange>,
97 ) -> Result<crate::ExprDict<Self::TargetU>, Self::Error> {
98 let crate::ExprDict {
99 range,
100 keys,
101 values,
102 } = node;
103 let context = self.will_map_user(&range);
104 assert_eq!(keys.len(), values.len());
105 let mut located_keys = Vec::with_capacity(keys.len());
106 let mut located_values = Vec::with_capacity(values.len());
107 for (key, value) in keys.into_iter().zip(values.into_iter()) {
108 located_keys.push(self.fold(key)?);
109 located_values.push(self.fold(value)?);
110 }
111 let range = self.map_user(range, context)?;
112 Ok(crate::ExprDict {
113 range,
114 keys: located_keys,
115 values: located_values,
116 })
117 }
118
119 fn fold_expr_if_exp(
120 &mut self,
121 node: crate::ExprIfExp<TextRange>,
122 ) -> Result<crate::ExprIfExp<Self::TargetU>, Self::Error> {
123 let crate::ExprIfExp {
124 range,
125 test,
126 body,
127 orelse,
128 } = node;
129 let context = self.will_map_user(&range);
130 let body = self.fold(body)?;
131 let test = self.fold(test)?;
132 let orelse = self.fold(orelse)?;
133 let range = self.map_user(range, context)?;
134 Ok(crate::ExprIfExp {
135 range,
136 test,
137 body,
138 orelse,
139 })
140 }
141
142 fn fold_stmt_class_def(
143 &mut self,
144 node: crate::StmtClassDef<TextRange>,
145 ) -> Result<crate::StmtClassDef<Self::TargetU>, Self::Error> {
146 let crate::StmtClassDef {
147 name,
148 bases,
149 keywords,
150 body,
151 decorator_list,
152 type_params,
153 range,
154 } = node;
155 let decorator_list = self.fold(decorator_list)?;
156 let context = self.will_map_user(&range);
157
158 let name = self.fold(name)?;
159 let type_params = self.fold(type_params)?;
160 let bases = self.fold(bases)?;
161 let keywords = self.fold(keywords)?;
162 let body = self.fold(body)?;
163 let range = self.map_user(range, context)?;
164
165 Ok(crate::StmtClassDef {
166 name,
167 bases,
168 keywords,
169 body,
170 decorator_list,
171 type_params,
172 range,
173 })
174 }
175 fn fold_stmt_function_def(
176 &mut self,
177 node: crate::StmtFunctionDef<TextRange>,
178 ) -> Result<crate::StmtFunctionDef<Self::TargetU>, Self::Error> {
179 let crate::StmtFunctionDef {
180 name,
181 args,
182 body,
183 decorator_list,
184 returns,
185 type_comment,
186 range,
187 type_params,
188 } = node;
189 let decorator_list = self.fold(decorator_list)?;
190 let context = self.will_map_user(&range);
191
192 let name = self.fold(name)?;
193 let type_params = self.fold(type_params)?;
194 let args: Box<crate::Arguments<SourceRange>> = self.fold(args)?;
195 let returns = self.fold(returns)?;
196 let body = self.fold(body)?;
197 let type_comment = self.fold(type_comment)?;
198 let range = self.map_user(range, context)?;
199 Ok(crate::StmtFunctionDef {
200 name,
201 args,
202 body,
203 decorator_list,
204 returns,
205 type_params,
206 type_comment,
207 range,
208 })
209 }
210 fn fold_stmt_async_function_def(
211 &mut self,
212 node: crate::StmtAsyncFunctionDef<TextRange>,
213 ) -> Result<crate::StmtAsyncFunctionDef<Self::TargetU>, Self::Error> {
214 let crate::StmtAsyncFunctionDef {
215 name,
216 args,
217 body,
218 decorator_list,
219 returns,
220 type_comment,
221 type_params,
222 range,
223 } = node;
224 let decorator_list = self.fold(decorator_list)?;
225 let context = self.will_map_user(&range);
226
227 let name = self.fold(name)?;
228 let type_params = self.fold(type_params)?;
229 let args: Box<crate::Arguments<SourceRange>> = self.fold(args)?;
230 let returns = self.fold(returns)?;
231 let body = self.fold(body)?;
232 let type_comment = self.fold(type_comment)?;
233 let range = self.map_user(range, context)?;
234 Ok(crate::StmtAsyncFunctionDef {
235 name,
236 args,
237 body,
238 decorator_list,
239 returns,
240 type_comment,
241 type_params,
242 range,
243 })
244 }
245 fn fold_expr_joined_str(
246 &mut self,
247 node: crate::ExprJoinedStr<TextRange>,
248 ) -> Result<crate::ExprJoinedStr<Self::TargetU>, Self::Error> {
249 let start = self.locate(node.range.start());
250 let end = self.locate_only(node.range.end());
251 let location = SourceRange::new(start, end);
252 linear_locate_expr_joined_str(self, node, location)
253 }
254
255 fn fold_expr_call(
256 &mut self,
257 node: crate::ExprCall<TextRange>,
258 ) -> Result<crate::ExprCall<Self::TargetU>, Self::Error> {
259 let crate::ExprCall {
260 range,
261 func,
262 args,
263 keywords,
264 } = node;
265 let context = self.will_map_user(&range);
266 let func = self.fold(func)?;
267 let keywords = LinearLookaheadLocator(self).fold(keywords)?;
268 let args = self.fold(args)?;
269 let range = self.map_user(range, context)?;
270 Ok(crate::ExprCall {
271 range,
272 func,
273 args,
274 keywords,
275 })
276 }
277}
278
279struct LinearLookaheadLocator<'a, 'b>(&'b mut LinearLocator<'a>);
280
281impl crate::fold::Fold<TextRange> for LinearLookaheadLocator<'_, '_> {
282 type TargetU = SourceRange;
283 type Error = std::convert::Infallible;
284 type UserContext = SourceLocation;
285
286 fn will_map_user(&mut self, user: &TextRange) -> Self::UserContext {
287 self.0.locate_only(user.start())
288 }
289
290 fn map_user(
291 &mut self,
292 user: TextRange,
293 start: Self::UserContext,
294 ) -> Result<Self::TargetU, Self::Error> {
295 let end = self.0.locate_only(user.end());
296 Ok((start..end).into())
297 }
298}