1use std::borrow::Cow;
2use std::path::Path;
3
4use rustc_hash::FxHashMap;
5
6use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer, indentation_at_offset};
7use ruff_source_file::LineRanges;
8use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
9
10use crate::name::{Name, QualifiedName, QualifiedNameBuilder};
11use crate::statement_visitor::StatementVisitor;
12use crate::token::Tokens;
13use crate::token::parenthesized_range;
14use crate::visitor::Visitor;
15use crate::{
16 self as ast, Arguments, AtomicNodeIndex, CmpOp, ConstantValue, DictItem, ExceptHandler, Expr,
17 ExprNoneLiteral, InterpolatedStringElement, MatchCase, Operator, Pattern, Stmt, TypeParam,
18};
19use crate::{AnyNodeRef, ExprContext};
20
21pub const fn is_compound_statement(stmt: &Stmt) -> bool {
23 matches!(
24 stmt,
25 Stmt::FunctionDef(_)
26 | Stmt::ClassDef(_)
27 | Stmt::While(_)
28 | Stmt::For(_)
29 | Stmt::Match(_)
30 | Stmt::With(_)
31 | Stmt::If(_)
32 | Stmt::Try(_)
33 )
34}
35
36fn is_iterable_initializer<F>(id: &str, is_builtin: F) -> bool
37where
38 F: Fn(&str) -> bool,
39{
40 matches!(id, "list" | "tuple" | "set" | "dict" | "frozenset") && is_builtin(id)
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum SideEffect {
47 Absent,
49 Possible,
52 Present,
54}
55
56impl SideEffect {
57 pub const fn is_present(self) -> bool {
58 matches!(self, Self::Present)
59 }
60
61 pub const fn is_absent(self) -> bool {
62 matches!(self, Self::Absent)
63 }
64
65 #[must_use]
66 pub const fn merge(self, other: Self) -> Self {
67 match (self, other) {
68 (Self::Present, _) | (_, Self::Present) => Self::Present,
69 (Self::Possible, _) | (_, Self::Possible) => Self::Possible,
70 _ => Self::Absent,
71 }
72 }
73
74 fn from_expr(expr: &Expr, is_builtin: &dyn Fn(&str) -> bool) -> Self {
76 match expr {
77 Expr::Call(ast::ExprCall {
79 func, arguments, ..
80 }) if arguments.is_empty() => {
81 if let Expr::Name(ast::ExprName { id, .. }) = func.as_ref() {
82 if is_iterable_initializer(id.as_str(), |id| is_builtin(id)) {
83 return Self::Absent;
84 }
85 }
86 Self::Present
87 }
88
89 Expr::BinOp(ast::ExprBinOp { left, right, .. }) => {
91 if is_known_safe_binop_operand(left) && is_known_safe_binop_operand(right) {
92 Self::Absent
93 } else {
94 Self::Present
95 }
96 }
97
98 Expr::FString(ast::ExprFString { value, .. }) => {
100 if value.elements().any(has_uncertain_interpolation) {
101 Self::Possible
102 } else {
103 Self::Absent
104 }
105 }
106 Expr::TString(ast::ExprTString { value, .. }) => {
107 if value.elements().any(has_uncertain_interpolation) {
108 Self::Possible
109 } else {
110 Self::Absent
111 }
112 }
113
114 Expr::Named(_) => Self::Present,
116
117 Expr::Await(_)
119 | Expr::Call(_)
120 | Expr::DictComp(_)
121 | Expr::Generator(_)
122 | Expr::ListComp(_)
123 | Expr::SetComp(_)
124 | Expr::Subscript(_)
125 | Expr::Yield(_)
126 | Expr::YieldFrom(_)
127 | Expr::IpyEscapeCommand(_) => Self::Present,
128
129 Expr::BoolOp(_)
131 | Expr::Compare(_)
132 | Expr::Dict(_)
133 | Expr::If(_)
134 | Expr::Lambda(_)
135 | Expr::List(_)
136 | Expr::Set(_)
137 | Expr::Slice(_)
138 | Expr::Starred(_)
139 | Expr::Tuple(_)
140 | Expr::UnaryOp(_)
141 | Expr::Attribute(_)
142 | Expr::Name(_)
143 | Expr::StringLiteral(_)
144 | Expr::BytesLiteral(_)
145 | Expr::NumberLiteral(_)
146 | Expr::Constant(_)
147 | Expr::BooleanLiteral(_)
148 | Expr::NoneLiteral(_)
149 | Expr::EllipsisLiteral(_) => Self::Absent,
150 }
151 }
152}
153
154const fn is_known_safe_binop_operand(expr: &Expr) -> bool {
155 match expr {
156 Expr::StringLiteral(_)
157 | Expr::BytesLiteral(_)
158 | Expr::NumberLiteral(_)
159 | Expr::Constant(_)
160 | Expr::BooleanLiteral(_)
161 | Expr::NoneLiteral(_)
162 | Expr::EllipsisLiteral(_)
163 | Expr::FString(_)
164 | Expr::List(_)
165 | Expr::Tuple(_)
166 | Expr::Set(_)
167 | Expr::Dict(_)
168 | Expr::ListComp(_)
169 | Expr::SetComp(_)
170 | Expr::DictComp(_) => true,
171
172 Expr::BoolOp(_)
173 | Expr::Named(_)
174 | Expr::BinOp(_)
175 | Expr::UnaryOp(_)
176 | Expr::Lambda(_)
177 | Expr::If(_)
178 | Expr::Compare(_)
179 | Expr::Call(_)
180 | Expr::Generator(_)
181 | Expr::Await(_)
182 | Expr::Yield(_)
183 | Expr::YieldFrom(_)
184 | Expr::Attribute(_)
185 | Expr::Subscript(_)
186 | Expr::Starred(_)
187 | Expr::Name(_)
188 | Expr::Slice(_)
189 | Expr::IpyEscapeCommand(_)
190 | Expr::TString(_) => false,
191 }
192}
193
194fn is_definitely_side_effect_free_interpolation_expr(expr: &Expr) -> bool {
195 matches!(
196 expr,
197 Expr::NumberLiteral(_)
198 | Expr::Constant(_)
199 | Expr::BooleanLiteral(_)
200 | Expr::NoneLiteral(_)
201 | Expr::EllipsisLiteral(_)
202 | Expr::StringLiteral(_)
203 | Expr::BytesLiteral(_)
204 )
205}
206
207fn has_uncertain_interpolation(element: &InterpolatedStringElement) -> bool {
208 match element {
209 InterpolatedStringElement::Literal(_) => false,
210 InterpolatedStringElement::Interpolation(interp) => {
211 !is_definitely_side_effect_free_interpolation_expr(&interp.expression)
212 || interp
213 .format_spec
214 .as_ref()
215 .is_some_and(|spec| spec.elements.iter().any(has_uncertain_interpolation))
216 }
217 }
218}
219
220pub fn contains_effect<F>(expr: &Expr, is_builtin: F) -> bool
225where
226 F: Fn(&str) -> bool,
227{
228 side_effect(expr, is_builtin).is_present()
229}
230
231pub fn side_effect<F>(expr: &Expr, is_builtin: F) -> SideEffect
239where
240 F: Fn(&str) -> bool,
241{
242 let mut effect = SideEffect::Absent;
243 any_over_expr(expr, |expr| {
244 match SideEffect::from_expr(expr, &is_builtin) {
245 SideEffect::Present => {
246 effect = SideEffect::Present;
247 true
248 }
249 SideEffect::Possible => {
250 effect = effect.merge(SideEffect::Possible);
251 false
252 }
253 SideEffect::Absent => false,
254 }
255 });
256 effect
257}
258
259pub fn any_over_expr<F>(expr: &Expr, mut func: F) -> bool
262where
263 F: FnMut(&Expr) -> bool,
264{
265 fn inner(expr: &Expr, func: &mut dyn FnMut(&Expr) -> bool) -> bool {
266 if func(expr) {
267 return true;
268 }
269 match expr {
270 Expr::BoolOp(ast::ExprBoolOp { values, .. }) => {
271 values.iter().any(|expr| any_over_expr(expr, &mut *func))
272 }
273 Expr::FString(ast::ExprFString { value, .. }) => value
274 .elements()
275 .any(|expr| any_over_interpolated_string_element(expr, &mut *func)),
276 Expr::TString(ast::ExprTString { value, .. }) => value
277 .elements()
278 .any(|expr| any_over_interpolated_string_element(expr, &mut *func)),
279 Expr::Named(ast::ExprNamed {
280 target,
281 value,
282 range: _,
283 node_index: _,
284 }) => any_over_expr(target, &mut *func) || any_over_expr(value, &mut *func),
285 Expr::BinOp(ast::ExprBinOp { left, right, .. }) => {
286 any_over_expr(left, &mut *func) || any_over_expr(right, &mut *func)
287 }
288 Expr::UnaryOp(ast::ExprUnaryOp { operand, .. }) => any_over_expr(operand, func),
289 Expr::Lambda(ast::ExprLambda {
290 body, parameters, ..
291 }) => {
292 parameters
293 .iter()
294 .flat_map(|parameters| parameters.iter_non_variadic_params())
295 .filter_map(|parameter| parameter.default.as_deref())
296 .any(|default| any_over_expr(default, &mut *func))
297 || any_over_expr(body, func)
298 }
299 Expr::If(ast::ExprIf {
300 test,
301 body,
302 orelse,
303 range: _,
304 node_index: _,
305 }) => {
306 any_over_expr(test, &mut *func)
307 || any_over_expr(body, &mut *func)
308 || any_over_expr(orelse, &mut *func)
309 }
310 Expr::Dict(ast::ExprDict { items, .. }) => {
311 items.iter().any(|ast::DictItem { key, value }| {
312 any_over_expr(value, &mut *func)
313 || key
314 .as_ref()
315 .is_some_and(|key| any_over_expr(key, &mut *func))
316 })
317 }
318 Expr::Set(ast::ExprSet { elts, .. })
319 | Expr::List(ast::ExprList { elts, .. })
320 | Expr::Tuple(ast::ExprTuple { elts, .. }) => {
321 elts.iter().any(|expr| any_over_expr(expr, &mut *func))
322 }
323 Expr::ListComp(ast::ExprListComp {
324 elt,
325 generators,
326 range: _,
327 node_index: _,
328 })
329 | Expr::SetComp(ast::ExprSetComp {
330 elt,
331 generators,
332 range: _,
333 node_index: _,
334 })
335 | Expr::Generator(ast::ExprGenerator {
336 elt,
337 generators,
338 range: _,
339 node_index: _,
340 parenthesized: _,
341 }) => {
342 any_over_expr(elt, &mut *func)
343 || generators.iter().any(|generator| {
344 any_over_expr(&generator.target, &mut *func)
345 || any_over_expr(&generator.iter, &mut *func)
346 || generator
347 .ifs
348 .iter()
349 .any(|expr| any_over_expr(expr, &mut *func))
350 })
351 }
352 Expr::DictComp(ast::ExprDictComp {
353 key,
354 value,
355 generators,
356 range: _,
357 node_index: _,
358 }) => {
359 key.as_deref()
360 .is_some_and(|key| any_over_expr(key, &mut *func))
361 || any_over_expr(value, &mut *func)
362 || generators.iter().any(|generator| {
363 any_over_expr(&generator.target, &mut *func)
364 || any_over_expr(&generator.iter, &mut *func)
365 || generator
366 .ifs
367 .iter()
368 .any(|expr| any_over_expr(expr, &mut *func))
369 })
370 }
371 Expr::Await(ast::ExprAwait {
372 value,
373 range: _,
374 node_index: _,
375 })
376 | Expr::YieldFrom(ast::ExprYieldFrom {
377 value,
378 range: _,
379 node_index: _,
380 })
381 | Expr::Attribute(ast::ExprAttribute { value, .. })
382 | Expr::Starred(ast::ExprStarred { value, .. }) => any_over_expr(value, func),
383 Expr::Yield(ast::ExprYield {
384 value,
385 range: _,
386 node_index: _,
387 }) => value
388 .as_ref()
389 .is_some_and(|value| any_over_expr(value, func)),
390 Expr::Compare(ast::ExprCompare {
391 left, comparators, ..
392 }) => {
393 any_over_expr(left, &mut *func)
394 || comparators
395 .iter()
396 .any(|expr| any_over_expr(expr, &mut *func))
397 }
398 Expr::Call(ast::ExprCall {
399 func: call_func,
400 arguments,
401 range_start: _,
402 node_index: _,
403 }) => {
404 any_over_expr(call_func, &mut *func)
407 || arguments
408 .args
409 .iter()
410 .any(|expr| any_over_expr(expr, &mut *func))
411 || arguments
412 .keywords
413 .iter()
414 .any(|keyword| any_over_expr(&keyword.value, &mut *func))
415 }
416 Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
417 any_over_expr(value, &mut *func) || any_over_expr(slice, &mut *func)
418 }
419 Expr::Slice(ast::ExprSlice {
420 lower,
421 upper,
422 step,
423 range: _,
424 node_index: _,
425 }) => {
426 lower
427 .as_ref()
428 .is_some_and(|value| any_over_expr(value, &mut *func))
429 || upper
430 .as_ref()
431 .is_some_and(|value| any_over_expr(value, &mut *func))
432 || step
433 .as_ref()
434 .is_some_and(|value| any_over_expr(value, &mut *func))
435 }
436 Expr::Name(_)
437 | Expr::StringLiteral(_)
438 | Expr::BytesLiteral(_)
439 | Expr::NumberLiteral(_)
440 | Expr::Constant(_)
441 | Expr::BooleanLiteral(_)
442 | Expr::NoneLiteral(_)
443 | Expr::EllipsisLiteral(_)
444 | Expr::IpyEscapeCommand(_) => false,
445 }
446 }
447
448 inner(expr, &mut func)
449}
450
451fn any_over_type_param(type_param: &TypeParam, func: &mut dyn FnMut(&Expr) -> bool) -> bool {
452 match type_param {
453 TypeParam::TypeVar(ast::TypeParamTypeVar { bound, default, .. }) => {
454 bound
455 .as_ref()
456 .is_some_and(|value| any_over_expr(value, &mut *func))
457 || default
458 .as_ref()
459 .is_some_and(|value| any_over_expr(value, &mut *func))
460 }
461 TypeParam::TypeVarTuple(ast::TypeParamTypeVarTuple { default, .. }) => default
462 .as_ref()
463 .is_some_and(|value| any_over_expr(value, &mut *func)),
464 TypeParam::ParamSpec(ast::TypeParamParamSpec { default, .. }) => default
465 .as_ref()
466 .is_some_and(|value| any_over_expr(value, &mut *func)),
467 }
468}
469
470fn any_over_pattern(pattern: &Pattern, func: &mut dyn FnMut(&Expr) -> bool) -> bool {
471 match pattern {
472 Pattern::MatchValue(ast::PatternMatchValue {
473 value,
474 range: _,
475 node_index: _,
476 }) => any_over_expr(value, func),
477 Pattern::MatchSingleton(_) => false,
478 Pattern::MatchSequence(ast::PatternMatchSequence { patterns, .. }) => patterns
479 .iter()
480 .any(|pattern| any_over_pattern(pattern, &mut *func)),
481 Pattern::MatchMapping(ast::PatternMatchMapping { keys, patterns, .. }) => {
482 keys.iter().any(|key| any_over_expr(key, &mut *func))
483 || patterns
484 .iter()
485 .any(|pattern| any_over_pattern(pattern, &mut *func))
486 }
487 Pattern::MatchClass(ast::PatternMatchClass { cls, arguments, .. }) => {
488 any_over_expr(cls, &mut *func)
489 || arguments
490 .patterns
491 .iter()
492 .any(|pattern| any_over_pattern(pattern, &mut *func))
493 || arguments
494 .keywords
495 .iter()
496 .any(|keyword| any_over_pattern(&keyword.pattern, &mut *func))
497 }
498 Pattern::MatchStar(_) => false,
499 Pattern::MatchAs(ast::PatternMatchAs { pattern, .. }) => pattern
500 .as_ref()
501 .is_some_and(|pattern| any_over_pattern(pattern, func)),
502 Pattern::MatchOr(ast::PatternMatchOr { patterns, .. }) => patterns
503 .iter()
504 .any(|pattern| any_over_pattern(pattern, &mut *func)),
505 }
506}
507
508fn any_over_interpolated_string_element(
509 element: &ast::InterpolatedStringElement,
510 func: &mut dyn FnMut(&Expr) -> bool,
511) -> bool {
512 match element {
513 ast::InterpolatedStringElement::Literal(_) => false,
514 ast::InterpolatedStringElement::Interpolation(ast::InterpolatedElement {
515 expression,
516 format_spec,
517 ..
518 }) => {
519 any_over_expr(expression, &mut *func)
520 || format_spec.as_ref().is_some_and(|spec| {
521 spec.elements.iter().any(|spec_element| {
522 any_over_interpolated_string_element(spec_element, &mut *func)
523 })
524 })
525 }
526 }
527}
528
529fn any_over_stmt<F>(stmt: &Stmt, mut func: F) -> bool
530where
531 F: FnMut(&Expr) -> bool,
532{
533 fn inner(stmt: &Stmt, func: &mut dyn FnMut(&Expr) -> bool) -> bool {
534 match stmt {
535 Stmt::FunctionDef(ast::StmtFunctionDef {
536 parameters,
537 type_params,
538 body,
539 decorator_list,
540 returns,
541 ..
542 }) => {
543 parameters.iter().any(|param| {
544 param
545 .default()
546 .is_some_and(|default| any_over_expr(default, &mut *func))
547 || param
548 .annotation()
549 .is_some_and(|annotation| any_over_expr(annotation, &mut *func))
550 }) || type_params.as_ref().is_some_and(|type_params| {
551 type_params
552 .iter()
553 .any(|type_param| any_over_type_param(type_param, &mut *func))
554 }) || body.iter().any(|stmt| any_over_stmt(stmt, &mut *func))
555 || decorator_list
556 .iter()
557 .any(|decorator| any_over_expr(&decorator.expression, &mut *func))
558 || returns
559 .as_ref()
560 .is_some_and(|value| any_over_expr(value, func))
561 }
562 Stmt::ClassDef(ast::StmtClassDef {
563 arguments,
564 type_params,
565 body,
566 decorator_list,
567 ..
568 }) => {
569 arguments
572 .as_deref()
573 .is_some_and(|Arguments { args, keywords, .. }| {
574 args.iter().any(|expr| any_over_expr(expr, &mut *func))
575 || keywords
576 .iter()
577 .any(|keyword| any_over_expr(&keyword.value, &mut *func))
578 })
579 || type_params.as_ref().is_some_and(|type_params| {
580 type_params
581 .iter()
582 .any(|type_param| any_over_type_param(type_param, &mut *func))
583 })
584 || body.iter().any(|stmt| any_over_stmt(stmt, &mut *func))
585 || decorator_list
586 .iter()
587 .any(|decorator| any_over_expr(&decorator.expression, &mut *func))
588 }
589 Stmt::Return(ast::StmtReturn {
590 value,
591 range: _,
592 node_index: _,
593 }) => value
594 .as_ref()
595 .is_some_and(|value| any_over_expr(value, func)),
596 Stmt::Delete(ast::StmtDelete { targets, .. }) => {
597 targets.iter().any(|expr| any_over_expr(expr, &mut *func))
598 }
599 Stmt::TypeAlias(ast::StmtTypeAlias {
600 name,
601 type_params,
602 value,
603 ..
604 }) => {
605 any_over_expr(name, &mut *func)
606 || type_params.as_ref().is_some_and(|type_params| {
607 type_params
608 .iter()
609 .any(|type_param| any_over_type_param(type_param, &mut *func))
610 })
611 || any_over_expr(value, func)
612 }
613 Stmt::Assign(ast::StmtAssign { targets, value, .. }) => {
614 targets.iter().any(|expr| any_over_expr(expr, &mut *func))
615 || any_over_expr(value, func)
616 }
617 Stmt::AugAssign(ast::StmtAugAssign { target, value, .. }) => {
618 any_over_expr(target, &mut *func) || any_over_expr(value, &mut *func)
619 }
620 Stmt::AnnAssign(ast::StmtAnnAssign {
621 target,
622 annotation,
623 value,
624 ..
625 }) => {
626 any_over_expr(target, &mut *func)
627 || any_over_expr(annotation, &mut *func)
628 || value
629 .as_ref()
630 .is_some_and(|value| any_over_expr(value, &mut *func))
631 }
632 Stmt::For(ast::StmtFor {
633 target,
634 iter,
635 body,
636 orelse,
637 ..
638 }) => {
639 any_over_expr(target, &mut *func)
640 || any_over_expr(iter, &mut *func)
641 || any_over_body(body, &mut *func)
642 || any_over_body(orelse, &mut *func)
643 }
644 Stmt::While(ast::StmtWhile {
645 test, body, orelse, ..
646 }) => {
647 any_over_expr(test, &mut *func)
648 || any_over_body(body, &mut *func)
649 || any_over_body(orelse, &mut *func)
650 }
651 Stmt::If(ast::StmtIf {
652 test,
653 body,
654 elif_else_clauses,
655 ..
656 }) => {
657 any_over_expr(test, &mut *func)
658 || any_over_body(body, &mut *func)
659 || elif_else_clauses.iter().any(|clause| {
660 clause
661 .test
662 .as_ref()
663 .is_some_and(|test| any_over_expr(test, &mut *func))
664 || any_over_body(&clause.body, &mut *func)
665 })
666 }
667 Stmt::With(ast::StmtWith { items, body, .. }) => {
668 items.iter().any(|with_item| {
669 any_over_expr(&with_item.context_expr, &mut *func)
670 || with_item
671 .optional_vars
672 .as_ref()
673 .is_some_and(|expr| any_over_expr(expr, &mut *func))
674 }) || any_over_body(body, &mut *func)
675 }
676 Stmt::Raise(ast::StmtRaise {
677 exc,
678 cause,
679 range: _,
680 node_index: _,
681 }) => {
682 exc.as_ref()
683 .is_some_and(|value| any_over_expr(value, &mut *func))
684 || cause
685 .as_ref()
686 .is_some_and(|value| any_over_expr(value, &mut *func))
687 }
688 Stmt::Try(ast::StmtTry {
689 body,
690 handlers,
691 orelse,
692 finalbody,
693 ..
694 }) => {
695 any_over_body(body, &mut *func)
696 || handlers.iter().any(|handler| {
697 let ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler {
698 type_,
699 body,
700 ..
701 }) = handler;
702 type_
703 .as_ref()
704 .is_some_and(|expr| any_over_expr(expr, &mut *func))
705 || any_over_body(body, &mut *func)
706 })
707 || any_over_body(orelse, &mut *func)
708 || any_over_body(finalbody, &mut *func)
709 }
710 Stmt::Assert(ast::StmtAssert {
711 test,
712 msg,
713 range: _,
714 node_index: _,
715 }) => {
716 any_over_expr(test, &mut *func)
717 || msg
718 .as_ref()
719 .is_some_and(|value| any_over_expr(value, &mut *func))
720 }
721 Stmt::Match(ast::StmtMatch {
722 subject,
723 cases,
724 range: _,
725 node_index: _,
726 }) => {
727 any_over_expr(subject, &mut *func)
728 || cases.iter().any(|case| {
729 let MatchCase {
730 pattern,
731 guard,
732 body,
733 ..
734 } = case;
735 any_over_pattern(pattern, &mut *func)
736 || guard
737 .as_ref()
738 .is_some_and(|expr| any_over_expr(expr, &mut *func))
739 || any_over_body(body, &mut *func)
740 })
741 }
742 Stmt::Import(_) => false,
743 Stmt::ImportFrom(_) => false,
744 Stmt::Global(_) => false,
745 Stmt::Nonlocal(_) => false,
746 Stmt::Expr(ast::StmtExpr {
747 value,
748 range: _,
749 node_index: _,
750 }) => any_over_expr(value, func),
751 Stmt::Pass(_) | Stmt::Break(_) | Stmt::Continue(_) => false,
752 Stmt::IpyEscapeCommand(_) => false,
753 }
754 }
755
756 inner(stmt, &mut func)
757}
758
759pub fn any_over_body<F>(body: &[Stmt], mut func: F) -> bool
760where
761 F: FnMut(&Expr) -> bool,
762{
763 body.iter().any(|stmt| any_over_stmt(stmt, &mut func))
764}
765
766pub fn is_dunder(id: &str) -> bool {
767 id.starts_with("__") && id.ends_with("__")
768}
769
770pub fn is_sunder(id: &str) -> bool {
774 id.starts_with('_') && id.ends_with('_') && !id.starts_with("__") && !id.ends_with("__")
775}
776
777pub fn is_assignment_to_a_dunder(stmt: &Stmt) -> bool {
779 match stmt {
782 Stmt::Assign(ast::StmtAssign { targets, .. }) => {
783 if let [Expr::Name(ast::ExprName { id, .. })] = targets.as_slice() {
784 is_dunder(id)
785 } else {
786 false
787 }
788 }
789 Stmt::AnnAssign(ast::StmtAnnAssign { target, .. }) => {
790 if let Expr::Name(ast::ExprName { id, .. }) = target.as_ref() {
791 is_dunder(id)
792 } else {
793 false
794 }
795 }
796 _ => false,
797 }
798}
799
800const fn is_singleton(expr: &Expr) -> bool {
803 matches!(
804 expr,
805 Expr::NoneLiteral(_)
806 | Expr::BooleanLiteral(_)
807 | Expr::EllipsisLiteral(_)
808 | Expr::Constant(ast::ExprConstant {
809 value: ConstantValue::None | ConstantValue::Boolean(_) | ConstantValue::Ellipsis,
810 ..
811 })
812 )
813}
814
815pub fn is_constant(expr: &Expr) -> bool {
817 if let Expr::Tuple(tuple) = expr {
818 tuple.iter().all(is_constant)
819 } else {
820 expr.is_literal_expr()
821 }
822}
823
824pub fn is_constant_non_singleton(expr: &Expr) -> bool {
826 is_constant(expr) && !is_singleton(expr)
827}
828
829pub const fn is_const_true(expr: &Expr) -> bool {
831 matches!(
832 expr,
833 Expr::BooleanLiteral(ast::ExprBooleanLiteral { value: true, .. }),
834 )
835}
836
837pub const fn is_const_false(expr: &Expr) -> bool {
839 matches!(
840 expr,
841 Expr::BooleanLiteral(ast::ExprBooleanLiteral { value: false, .. }),
842 )
843}
844
845pub const fn is_mutable_iterable_initializer(expr: &Expr) -> bool {
847 matches!(
848 expr,
849 Expr::Set(_)
850 | Expr::SetComp(_)
851 | Expr::List(_)
852 | Expr::ListComp(_)
853 | Expr::Dict(_)
854 | Expr::DictComp(_)
855 )
856}
857
858pub fn extract_handled_exceptions(handlers: &[ExceptHandler]) -> Vec<&Expr> {
860 let mut handled_exceptions = Vec::new();
861 for handler in handlers {
862 match handler {
863 ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler { type_, .. }) => {
864 if let Some(type_) = type_ {
865 if let Expr::Tuple(tuple) = &**type_ {
866 for type_ in tuple {
867 handled_exceptions.push(type_);
868 }
869 } else {
870 handled_exceptions.push(type_);
871 }
872 }
873 }
874 }
875 }
876 handled_exceptions
877}
878
879pub fn map_callable(decorator: &Expr) -> &Expr {
883 if let Expr::Call(ast::ExprCall { func, .. }) = decorator {
884 func
886 } else {
887 decorator
889 }
890}
891
892pub fn map_subscript(expr: &Expr) -> &Expr {
895 if let Expr::Subscript(ast::ExprSubscript { value, .. }) = expr {
896 value
898 } else {
899 expr
901 }
902}
903
904pub fn map_starred(expr: &Expr) -> &Expr {
906 if let Expr::Starred(ast::ExprStarred { value, .. }) = expr {
907 value
909 } else {
910 expr
912 }
913}
914
915pub fn uses_magic_variable_access<F>(body: &[Stmt], is_builtin: F) -> bool
919where
920 F: Fn(&str) -> bool,
921{
922 any_over_body(body, |expr| {
923 if let Expr::Call(ast::ExprCall { func, .. }) = expr {
924 if let Expr::Name(ast::ExprName { id, .. }) = func.as_ref() {
925 if matches!(id.as_str(), "locals" | "globals" | "vars" | "exec" | "eval") {
926 if is_builtin(id.as_str()) {
927 return true;
928 }
929 }
930 }
931 }
932 false
933 })
934}
935
936pub fn format_import_from(level: u32, module: Option<&str>) -> Cow<'_, str> {
948 match (level, module) {
949 (0, Some(module)) => Cow::Borrowed(module),
950 (level, module) => {
951 let mut module_name =
952 String::with_capacity((level as usize) + module.map_or(0, str::len));
953 for _ in 0..level {
954 module_name.push('.');
955 }
956 if let Some(module) = module {
957 module_name.push_str(module);
958 }
959 Cow::Owned(module_name)
960 }
961 }
962}
963
964pub fn format_import_from_member(level: u32, module: Option<&str>, member: &str) -> String {
976 let mut qualified_name =
977 String::with_capacity((level as usize) + module.map_or(0, str::len) + 1 + member.len());
978 if level > 0 {
979 for _ in 0..level {
980 qualified_name.push('.');
981 }
982 }
983 if let Some(module) = module {
984 qualified_name.push_str(module);
985 qualified_name.push('.');
986 }
987 qualified_name.push_str(member);
988 qualified_name
989}
990
991pub fn to_module_path(package: &Path, path: &Path) -> Option<Vec<String>> {
996 path.strip_prefix(package.parent()?)
997 .ok()?
998 .iter()
999 .map(Path::new)
1000 .map(Path::file_stem)
1001 .map(|path| path.and_then(|path| path.to_os_string().into_string().ok()))
1002 .collect::<Option<Vec<String>>>()
1003}
1004
1005pub fn collect_import_from_member<'a>(
1017 level: u32,
1018 module: Option<&'a str>,
1019 member: &'a str,
1020) -> QualifiedName<'a> {
1021 let mut qualified_name_builder = QualifiedNameBuilder::with_capacity(
1022 level as usize
1023 + module
1024 .map(|module| module.split('.').count())
1025 .unwrap_or_default()
1026 + 1,
1027 );
1028
1029 if level > 0 {
1031 for _ in 0..level {
1032 qualified_name_builder.push(".");
1033 }
1034 }
1035
1036 if let Some(module) = module {
1038 qualified_name_builder.extend(module.split('.'));
1039 }
1040
1041 qualified_name_builder.push(member);
1043
1044 qualified_name_builder.build()
1045}
1046
1047pub fn from_relative_import<'a>(
1050 module: &'a [String],
1052 import: &[&'a str],
1054 tail: &[&'a str],
1056) -> Option<QualifiedName<'a>> {
1057 let mut qualified_name_builder =
1058 QualifiedNameBuilder::with_capacity(module.len() + import.len() + tail.len());
1059
1060 qualified_name_builder.extend(module.iter().map(String::as_str));
1062
1063 for segment in import {
1065 if *segment == "." {
1066 if qualified_name_builder.is_empty() {
1067 return None;
1068 }
1069 qualified_name_builder.pop();
1070 } else {
1071 qualified_name_builder.push(segment);
1072 }
1073 }
1074
1075 qualified_name_builder.extend_from_slice(tail);
1077
1078 Some(qualified_name_builder.build())
1079}
1080
1081pub fn resolve_imported_module_path<'a>(
1084 level: u32,
1085 module: Option<&'a str>,
1086 module_path: Option<&[String]>,
1087) -> Option<Cow<'a, str>> {
1088 if level == 0 {
1089 return Some(Cow::Borrowed(module.unwrap_or("")));
1090 }
1091
1092 let module_path = module_path?;
1093
1094 if level as usize >= module_path.len() {
1095 return None;
1096 }
1097
1098 let mut qualified_path = module_path[..module_path.len() - level as usize].join(".");
1099 if let Some(module) = module {
1100 if !qualified_path.is_empty() {
1101 qualified_path.push('.');
1102 }
1103 qualified_path.push_str(module);
1104 }
1105 Some(Cow::Owned(qualified_path))
1106}
1107
1108#[derive(Debug, Default)]
1110pub struct NameFinder<'a> {
1111 pub names: FxHashMap<&'a str, &'a ast::ExprName>,
1113}
1114
1115impl<'a> Visitor<'a> for NameFinder<'a> {
1116 fn visit_expr(&mut self, expr: &'a Expr) {
1117 if let Expr::Name(name) = expr {
1118 self.names.insert(&name.id, name);
1119 }
1120 crate::visitor::walk_expr(self, expr);
1121 }
1122}
1123
1124#[derive(Debug, Default)]
1126pub struct StoredNameFinder<'a> {
1127 pub names: FxHashMap<&'a str, &'a ast::ExprName>,
1129}
1130
1131impl<'a> Visitor<'a> for StoredNameFinder<'a> {
1132 fn visit_expr(&mut self, expr: &'a Expr) {
1133 if let Expr::Name(name) = expr {
1134 if name.ctx.is_store() {
1135 self.names.insert(&name.id, name);
1136 }
1137 }
1138 crate::visitor::walk_expr(self, expr);
1139 }
1140}
1141
1142#[derive(Default)]
1144pub struct ReturnStatementVisitor<'a> {
1145 pub returns: Vec<&'a ast::StmtReturn>,
1146 pub is_generator: bool,
1147}
1148
1149impl<'a> Visitor<'a> for ReturnStatementVisitor<'a> {
1150 fn visit_stmt(&mut self, stmt: &'a Stmt) {
1151 match stmt {
1152 Stmt::FunctionDef(_) | Stmt::ClassDef(_) => {
1153 }
1155 Stmt::Return(stmt) => self.returns.push(stmt),
1156 _ => crate::visitor::walk_stmt(self, stmt),
1157 }
1158 }
1159
1160 fn visit_expr(&mut self, expr: &'a Expr) {
1161 if let Expr::Yield(_) | Expr::YieldFrom(_) = expr {
1162 self.is_generator = true;
1163 } else {
1164 crate::visitor::walk_expr(self, expr);
1165 }
1166 }
1167}
1168
1169#[derive(Default)]
1171pub struct RaiseStatementVisitor<'a> {
1172 pub raises: Vec<(TextRange, Option<&'a Expr>, Option<&'a Expr>)>,
1173}
1174
1175impl<'a> StatementVisitor<'a> for RaiseStatementVisitor<'a> {
1176 fn visit_stmt(&mut self, stmt: &'a Stmt) {
1177 match stmt {
1178 Stmt::Raise(ast::StmtRaise {
1179 exc,
1180 cause,
1181 range: _,
1182 node_index: _,
1183 }) => {
1184 self.raises
1185 .push((stmt.range(), exc.as_deref(), cause.as_deref()));
1186 }
1187 Stmt::ClassDef(_) | Stmt::FunctionDef(_) | Stmt::Try(_) => {}
1188 Stmt::If(ast::StmtIf {
1189 body,
1190 elif_else_clauses,
1191 ..
1192 }) => {
1193 crate::statement_visitor::walk_body(self, body);
1194 for clause in elif_else_clauses {
1195 self.visit_elif_else_clause(clause);
1196 }
1197 }
1198 Stmt::While(ast::StmtWhile { body, .. })
1199 | Stmt::With(ast::StmtWith { body, .. })
1200 | Stmt::For(ast::StmtFor { body, .. }) => {
1201 crate::statement_visitor::walk_body(self, body);
1202 }
1203 Stmt::Match(ast::StmtMatch { cases, .. }) => {
1204 for case in cases {
1205 crate::statement_visitor::walk_body(self, &case.body);
1206 }
1207 }
1208 _ => {}
1209 }
1210 }
1211}
1212
1213#[derive(Debug, Default)]
1215pub struct AwaitVisitor {
1216 pub seen_await: bool,
1217}
1218
1219impl Visitor<'_> for AwaitVisitor {
1220 fn visit_stmt(&mut self, stmt: &Stmt) {
1221 match stmt {
1222 Stmt::FunctionDef(_) | Stmt::ClassDef(_) => (),
1223 Stmt::With(ast::StmtWith { is_async: true, .. }) => {
1224 self.seen_await = true;
1225 }
1226 Stmt::For(ast::StmtFor { is_async: true, .. }) => {
1227 self.seen_await = true;
1228 }
1229 _ => crate::visitor::walk_stmt(self, stmt),
1230 }
1231 }
1232
1233 fn visit_expr(&mut self, expr: &Expr) {
1234 if let Expr::Await(ast::ExprAwait { .. }) = expr {
1235 self.seen_await = true;
1236 } else {
1237 crate::visitor::walk_expr(self, expr);
1238 }
1239 }
1240
1241 fn visit_comprehension(&mut self, comprehension: &'_ crate::Comprehension) {
1242 if comprehension.is_async {
1243 self.seen_await = true;
1244 } else {
1245 crate::visitor::walk_comprehension(self, comprehension);
1246 }
1247 }
1248}
1249
1250pub fn is_docstring_stmt(stmt: &Stmt) -> bool {
1252 if let Stmt::Expr(ast::StmtExpr {
1253 value,
1254 range: _,
1255 node_index: _,
1256 }) = stmt
1257 {
1258 value.is_string_literal_expr()
1259 } else {
1260 false
1261 }
1262}
1263
1264pub fn is_stub_body(body: &[Stmt]) -> bool {
1268 !body.is_empty()
1269 && body.iter().all(|stmt| match stmt {
1270 Stmt::Pass(_) => true,
1271 Stmt::Expr(ast::StmtExpr { value, .. }) => value.is_ellipsis_literal_expr(),
1272 _ => false,
1273 })
1274}
1275
1276pub fn body_without_leading_docstring(body: &[Stmt]) -> &[Stmt] {
1278 match body.split_first() {
1279 Some((first, rest)) if is_docstring_stmt(first) => rest,
1280 _ => body,
1281 }
1282}
1283
1284pub fn on_conditional_branch<'a>(parents: &mut impl Iterator<Item = &'a Stmt>) -> bool {
1286 parents.any(|parent| {
1287 if matches!(parent, Stmt::If(_) | Stmt::While(_) | Stmt::Match(_)) {
1288 return true;
1289 }
1290 if let Stmt::Expr(ast::StmtExpr {
1291 value,
1292 range: _,
1293 node_index: _,
1294 }) = parent
1295 {
1296 if value.is_if_expr() {
1297 return true;
1298 }
1299 }
1300 false
1301 })
1302}
1303
1304pub fn in_nested_block<'a>(mut parents: impl Iterator<Item = &'a Stmt>) -> bool {
1306 parents.any(|parent| {
1307 matches!(
1308 parent,
1309 Stmt::Try(_) | Stmt::If(_) | Stmt::With(_) | Stmt::Match(_)
1310 )
1311 })
1312}
1313
1314pub fn is_unpacking_assignment(parent: &Stmt, child: &Expr) -> bool {
1316 match parent {
1317 Stmt::With(ast::StmtWith { items, .. }) => items.iter().any(|item| {
1318 if let Some(optional_vars) = &item.optional_vars {
1319 if optional_vars.is_tuple_expr() {
1320 if any_over_expr(optional_vars, |expr| expr == child) {
1321 return true;
1322 }
1323 }
1324 }
1325 false
1326 }),
1327 Stmt::Assign(ast::StmtAssign { targets, value, .. }) => {
1328 let value_is_tuple = matches!(
1330 value.as_ref(),
1331 Expr::Set(_) | Expr::List(_) | Expr::Tuple(_)
1332 );
1333 let targets_are_tuples = targets
1337 .iter()
1338 .all(|item| matches!(item, Expr::Set(_) | Expr::List(_) | Expr::Tuple(_)));
1339 let child_in_tuple = targets_are_tuples
1342 || targets.iter().any(|item| {
1343 matches!(item, Expr::Set(_) | Expr::List(_) | Expr::Tuple(_))
1344 && any_over_expr(item, |expr| expr == child)
1345 });
1346
1347 if child_in_tuple && !value_is_tuple {
1350 return true;
1351 }
1352
1353 if !child_in_tuple && value_is_tuple {
1356 return false;
1357 }
1358
1359 if child_in_tuple && value_is_tuple {
1365 return !targets_are_tuples;
1366 }
1367
1368 false
1369 }
1370 _ => false,
1371 }
1372}
1373
1374#[derive(Copy, Clone, Debug, PartialEq, is_macro::Is)]
1375pub enum Truthiness {
1376 True,
1378 False,
1380 Falsey,
1382 Truthy,
1384 None,
1386 Unknown,
1388}
1389
1390impl Truthiness {
1391 pub fn from_expr<F>(expr: &Expr, is_builtin: F) -> Self
1393 where
1394 F: Fn(&str) -> bool,
1395 {
1396 match expr {
1397 Expr::Lambda(_) => Self::Truthy,
1398 Expr::Generator(_) => Self::Truthy,
1399 Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => {
1400 if value.is_empty() {
1401 Self::Falsey
1402 } else {
1403 Self::Truthy
1404 }
1405 }
1406 Expr::BytesLiteral(ast::ExprBytesLiteral { value, .. }) => {
1407 if value.is_empty() {
1408 Self::Falsey
1409 } else {
1410 Self::Truthy
1411 }
1412 }
1413 Expr::NumberLiteral(ast::ExprNumberLiteral { value, .. }) => match value {
1414 ast::Number::Int(int) => {
1415 if *int == 0 {
1416 Self::Falsey
1417 } else {
1418 Self::Truthy
1419 }
1420 }
1421 ast::Number::Float(float) => {
1422 if *float == 0.0 {
1423 Self::Falsey
1424 } else {
1425 Self::Truthy
1426 }
1427 }
1428 ast::Number::Complex { real, imag, .. } => {
1429 if *real == 0.0 && *imag == 0.0 {
1430 Self::Falsey
1431 } else {
1432 Self::Truthy
1433 }
1434 }
1435 },
1436 Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => {
1437 if *value {
1438 Self::True
1439 } else {
1440 Self::False
1441 }
1442 }
1443 Expr::NoneLiteral(_) => Self::None,
1444 Expr::EllipsisLiteral(_) => Self::Truthy,
1445 Expr::FString(f_string) => {
1446 if is_empty_f_string(f_string) {
1447 Self::Falsey
1448 } else if is_non_empty_f_string(f_string) {
1449 Self::Truthy
1450 } else {
1451 Self::Unknown
1452 }
1453 }
1454 Expr::TString(_) => Self::Truthy,
1455 Expr::List(ast::ExprList { elts, .. })
1456 | Expr::Set(ast::ExprSet { elts, .. })
1457 | Expr::Tuple(ast::ExprTuple { elts, .. }) => {
1458 if elts.is_empty() {
1459 return Self::Falsey;
1460 }
1461
1462 if elts.iter().all(Expr::is_starred_expr) {
1463 Self::Unknown
1465 } else {
1466 Self::Truthy
1467 }
1468 }
1469 Expr::Dict(dict) => {
1470 if dict.is_empty() {
1471 return Self::Falsey;
1472 }
1473
1474 if dict
1478 .items
1479 .iter()
1480 .all(|item| matches!(item, DictItem { key: None, .. }))
1481 {
1482 Self::Unknown
1484 } else {
1485 Self::Truthy
1486 }
1487 }
1488 Expr::Call(ast::ExprCall {
1489 func, arguments, ..
1490 }) => {
1491 if let Expr::Name(ast::ExprName { id, .. }) = func.as_ref() {
1492 if is_iterable_initializer(id.as_str(), |id| is_builtin(id)) {
1493 if arguments.is_empty() {
1494 Self::Falsey
1496 } else if let [argument] = &*arguments.args
1497 && arguments.keywords.is_empty()
1498 {
1499 match argument {
1501 Expr::NumberLiteral(_)
1506 | Expr::BooleanLiteral(_)
1507 | Expr::NoneLiteral(_)
1508 | Expr::EllipsisLiteral(_)
1509 | Expr::TString(_)
1510 | Expr::Lambda(_)
1511 | Expr::Generator(_) => Self::Unknown,
1512 _ => Self::from_expr(argument, is_builtin),
1516 }
1517 } else {
1518 Self::Unknown
1519 }
1520 } else {
1521 Self::Unknown
1522 }
1523 } else {
1524 Self::Unknown
1525 }
1526 }
1527 _ => Self::Unknown,
1528 }
1529 }
1530
1531 pub fn into_bool(self) -> Option<bool> {
1532 match self {
1533 Self::True | Self::Truthy => Some(true),
1534 Self::False | Self::Falsey => Some(false),
1535 Self::None => Some(false),
1536 Self::Unknown => None,
1537 }
1538 }
1539}
1540
1541fn is_non_empty_f_string(expr: &ast::ExprFString) -> bool {
1544 fn inner(expr: &Expr) -> bool {
1545 match expr {
1546 Expr::Lambda(_) => true,
1548 Expr::Dict(_) => true,
1549 Expr::Set(_) => true,
1550 Expr::ListComp(_) => true,
1551 Expr::SetComp(_) => true,
1552 Expr::DictComp(_) => true,
1553 Expr::NumberLiteral(_) => true,
1554 Expr::Constant(_) => true,
1555 Expr::BooleanLiteral(_) => true,
1556 Expr::NoneLiteral(_) => true,
1557 Expr::EllipsisLiteral(_) => true,
1558 Expr::List(_) => true,
1559 Expr::Tuple(_) => true,
1560 Expr::TString(_) => true,
1561
1562 Expr::If(ast::ExprIf { body, orelse, .. }) => inner(body) && inner(orelse),
1564 Expr::Named(ast::ExprNamed { value, .. }) => inner(value),
1565
1566 Expr::BoolOp(ast::ExprBoolOp { .. }) => false,
1568 Expr::BinOp(ast::ExprBinOp { .. }) => false,
1569 Expr::UnaryOp(ast::ExprUnaryOp { .. }) => false,
1570 Expr::Compare(_) => false,
1572 Expr::Generator(_) => false,
1573 Expr::Await(_) => false,
1574 Expr::Yield(_) => false,
1575 Expr::YieldFrom(_) => false,
1576 Expr::Call(_) => false,
1577 Expr::Attribute(_) => false,
1578 Expr::Subscript(_) => false,
1579 Expr::Starred(_) => false,
1580 Expr::Name(_) => false,
1581 Expr::Slice(_) => false,
1582 Expr::IpyEscapeCommand(_) => false,
1583
1584 Expr::FString(f_string) => is_non_empty_f_string(f_string),
1586 Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => !value.is_empty(),
1588 Expr::BytesLiteral(_) => true,
1591 }
1592 }
1593
1594 expr.value.iter().any(|part| match part {
1595 ast::FStringPart::Literal(string_literal) => !string_literal.is_empty(),
1596 ast::FStringPart::FString(f_string) => {
1597 f_string.elements.iter().any(|element| match element {
1599 InterpolatedStringElement::Literal(string_literal) => !string_literal.is_empty(),
1600 InterpolatedStringElement::Interpolation(f_string) => {
1601 f_string.debug_text.is_some()
1602 || (f_string.format_spec.is_none() && inner(&f_string.expression))
1603 }
1604 })
1605 }
1606 })
1607}
1608
1609pub fn is_empty_f_string(expr: &ast::ExprFString) -> bool {
1612 fn inner(expr: &Expr) -> bool {
1613 match expr {
1614 Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => value.is_empty(),
1615 Expr::BytesLiteral(_) => false,
1619 Expr::FString(ast::ExprFString { value, .. }) => {
1620 is_empty_interpolated_elements(value.elements())
1621 }
1622 _ => false,
1623 }
1624 }
1625
1626 fn is_empty_interpolated_elements<'a>(
1627 mut elements: impl Iterator<Item = &'a InterpolatedStringElement>,
1628 ) -> bool {
1629 elements.all(|element| match element {
1630 InterpolatedStringElement::Literal(ast::InterpolatedStringLiteralElement {
1631 value,
1632 ..
1633 }) => value.is_empty(),
1634 InterpolatedStringElement::Interpolation(f_string) => {
1635 f_string.debug_text.is_none()
1636 && f_string.conversion.is_none()
1637 && f_string.format_spec.is_none()
1638 && inner(&f_string.expression)
1639 }
1640 })
1641 }
1642
1643 expr.value.iter().all(|part| match part {
1644 ast::FStringPart::Literal(string_literal) => string_literal.is_empty(),
1645 ast::FStringPart::FString(f_string) => {
1646 is_empty_interpolated_elements(f_string.elements.iter())
1647 }
1648 })
1649}
1650
1651pub fn generate_comparison(
1652 left: &Expr,
1653 ops: &[CmpOp],
1654 comparators: &[Expr],
1655 parent: AnyNodeRef,
1656 tokens: &Tokens,
1657 source: &str,
1658) -> String {
1659 let start = left.start();
1660 let end = comparators.last().map_or_else(|| left.end(), Ranged::end);
1661 let mut contents = String::with_capacity(usize::from(end - start));
1662
1663 contents.push_str(
1665 &source[parenthesized_range(left.into(), parent, tokens).unwrap_or(left.range())],
1666 );
1667
1668 for (op, comparator) in ops.iter().zip(comparators) {
1669 contents.push_str(match op {
1671 CmpOp::Eq => " == ",
1672 CmpOp::NotEq => " != ",
1673 CmpOp::Lt => " < ",
1674 CmpOp::LtE => " <= ",
1675 CmpOp::Gt => " > ",
1676 CmpOp::GtE => " >= ",
1677 CmpOp::In => " in ",
1678 CmpOp::NotIn => " not in ",
1679 CmpOp::Is => " is ",
1680 CmpOp::IsNot => " is not ",
1681 });
1682
1683 contents.push_str(
1685 &source[parenthesized_range(comparator.into(), parent, tokens)
1686 .unwrap_or(comparator.range())],
1687 );
1688 }
1689
1690 contents
1691}
1692
1693pub fn pep_604_optional(expr: &Expr) -> Expr {
1695 ast::ExprBinOp {
1696 left: Box::new(expr.clone()),
1697 op: Operator::BitOr,
1698 right: Box::new(Expr::NoneLiteral(ExprNoneLiteral::default())),
1699 range: TextRange::default(),
1700 node_index: AtomicNodeIndex::NONE,
1701 }
1702 .into()
1703}
1704
1705pub fn pep_604_union(elts: &[Expr]) -> Expr {
1707 match elts {
1708 [] => Expr::Tuple(ast::ExprTuple {
1709 elts: vec![],
1710 ctx: ExprContext::Load,
1711 range: TextRange::default(),
1712 node_index: AtomicNodeIndex::NONE,
1713 parenthesized: true,
1714 runtime_elts: None,
1715 }),
1716 [Expr::Tuple(ast::ExprTuple { elts, .. })] => pep_604_union(elts),
1717 [elt] => elt.clone(),
1718 [rest @ .., elt] => Expr::BinOp(ast::ExprBinOp {
1719 left: Box::new(pep_604_union(rest)),
1720 op: Operator::BitOr,
1721 right: Box::new(pep_604_union(std::slice::from_ref(elt))),
1722 range: TextRange::default(),
1723 node_index: AtomicNodeIndex::NONE,
1724 }),
1725 }
1726}
1727
1728pub fn typing_optional(elt: Expr, binding: Name) -> Expr {
1730 Expr::Subscript(ast::ExprSubscript {
1731 value: Box::new(Expr::Name(ast::ExprName {
1732 id: binding,
1733 range: TextRange::default(),
1734 node_index: AtomicNodeIndex::NONE,
1735 ctx: ExprContext::Load,
1736 })),
1737 slice: Box::new(elt),
1738 ctx: ExprContext::Load,
1739 range: TextRange::default(),
1740 node_index: AtomicNodeIndex::NONE,
1741 })
1742}
1743
1744pub fn typing_union(elts: &[Expr], binding: Name) -> Expr {
1749 Expr::Subscript(ast::ExprSubscript {
1750 value: Box::new(Expr::Name(ast::ExprName {
1751 id: binding,
1752 range: TextRange::default(),
1753 node_index: AtomicNodeIndex::NONE,
1754 ctx: ExprContext::Load,
1755 })),
1756 slice: Box::new(Expr::Tuple(ast::ExprTuple {
1757 range: TextRange::default(),
1758 node_index: AtomicNodeIndex::NONE,
1759 elts: elts.to_vec(),
1760 ctx: ExprContext::Load,
1761 parenthesized: false,
1762 runtime_elts: None,
1763 })),
1764 ctx: ExprContext::Load,
1765 range: TextRange::default(),
1766 node_index: AtomicNodeIndex::NONE,
1767 })
1768}
1769
1770pub fn comment_indentation_after(
1823 preceding: AnyNodeRef,
1824 comment_range: TextRange,
1825 source: &str,
1826) -> TextSize {
1827 let tokenizer = SimpleTokenizer::new(
1828 source,
1829 TextRange::new(source.full_line_end(preceding.end()), comment_range.end()),
1830 );
1831
1832 tokenizer
1833 .filter_map(|token| {
1834 if token.kind() == SimpleTokenKind::Comment {
1835 indentation_at_offset(token.start(), source).map(TextLen::text_len)
1836 } else {
1837 None
1838 }
1839 })
1840 .min()
1841 .unwrap_or_default()
1842}
1843
1844pub fn is_dotted_name(expr: &ast::Expr) -> bool {
1845 match expr {
1846 ast::Expr::Name(_) => true,
1847 ast::Expr::Attribute(ast::ExprAttribute { value, .. }) => is_dotted_name(value),
1848 _ => false,
1849 }
1850}
1851
1852#[cfg(test)]
1853mod tests {
1854 use std::borrow::Cow;
1855 use std::cell::RefCell;
1856 use std::vec;
1857
1858 use ruff_text_size::TextRange;
1859
1860 use crate::helpers::{any_over_stmt, any_over_type_param, resolve_imported_module_path};
1861 use crate::{
1862 AtomicNodeIndex, Expr, ExprContext, ExprName, ExprNumberLiteral, Identifier, Int, Number,
1863 Stmt, StmtTypeAlias, TypeParam, TypeParamParamSpec, TypeParamTypeVar,
1864 TypeParamTypeVarTuple, TypeParams,
1865 };
1866
1867 #[test]
1868 fn resolve_import() {
1869 assert_eq!(
1871 resolve_imported_module_path(0, Some("foo"), None),
1872 Some(Cow::Borrowed("foo"))
1873 );
1874
1875 assert_eq!(
1877 resolve_imported_module_path(
1878 1,
1879 Some("foo"),
1880 Some(&["bar".to_string(), "baz".to_string()])
1881 ),
1882 Some(Cow::Owned("bar.foo".to_string()))
1883 );
1884
1885 assert_eq!(resolve_imported_module_path(1, Some("foo"), None), None);
1888
1889 assert_eq!(
1892 resolve_imported_module_path(1, Some("foo"), Some(&["bar".to_string()])),
1893 None,
1894 );
1895 assert_eq!(
1896 resolve_imported_module_path(2, Some("foo"), Some(&["bar".to_string()])),
1897 None
1898 );
1899 }
1900
1901 #[test]
1902 fn any_over_stmt_type_alias() {
1903 let seen = RefCell::new(Vec::new());
1904 let name = Expr::Name(ExprName {
1905 id: "x".into(),
1906 range: TextRange::default(),
1907 node_index: AtomicNodeIndex::NONE,
1908 ctx: ExprContext::Load,
1909 });
1910 let constant_one = Expr::NumberLiteral(ExprNumberLiteral {
1911 value: Number::Int(Int::from(1u8)),
1912 range: TextRange::default(),
1913 node_index: AtomicNodeIndex::NONE,
1914 });
1915 let constant_two = Expr::NumberLiteral(ExprNumberLiteral {
1916 value: Number::Int(Int::from(2u8)),
1917 range: TextRange::default(),
1918 node_index: AtomicNodeIndex::NONE,
1919 });
1920 let constant_three = Expr::NumberLiteral(ExprNumberLiteral {
1921 value: Number::Int(Int::from(3u8)),
1922 range: TextRange::default(),
1923 node_index: AtomicNodeIndex::NONE,
1924 });
1925 let type_var_one = TypeParam::TypeVar(TypeParamTypeVar {
1926 range: TextRange::default(),
1927 node_index: AtomicNodeIndex::NONE,
1928 bound: Some(Box::new(constant_one.clone())),
1929 default: None,
1930 name: Identifier::new("x", TextRange::default()),
1931 });
1932 let type_var_two = TypeParam::TypeVar(TypeParamTypeVar {
1933 range: TextRange::default(),
1934 node_index: AtomicNodeIndex::NONE,
1935 bound: None,
1936 default: Some(Box::new(constant_two.clone())),
1937 name: Identifier::new("x", TextRange::default()),
1938 });
1939 let type_alias = Stmt::TypeAlias(StmtTypeAlias {
1940 name: Box::new(name.clone()),
1941 type_params: Some(Box::new(TypeParams {
1942 type_params: vec![type_var_one, type_var_two],
1943 range: TextRange::default(),
1944 node_index: AtomicNodeIndex::NONE,
1945 runtime_type_params: None,
1946 })),
1947 value: Box::new(constant_three.clone()),
1948 range: TextRange::default(),
1949 node_index: AtomicNodeIndex::NONE,
1950 });
1951 assert!(!any_over_stmt(&type_alias, |expr| {
1952 seen.borrow_mut().push(expr.clone());
1953 false
1954 }));
1955 assert_eq!(
1956 seen.take(),
1957 vec![name, constant_one, constant_two, constant_three]
1958 );
1959 }
1960
1961 #[test]
1962 fn any_over_type_param_type_var() {
1963 let type_var_no_bound = TypeParam::TypeVar(TypeParamTypeVar {
1964 range: TextRange::default(),
1965 node_index: AtomicNodeIndex::NONE,
1966 bound: None,
1967 default: None,
1968 name: Identifier::new("x", TextRange::default()),
1969 });
1970 assert!(!any_over_type_param(&type_var_no_bound, &mut |_expr| true));
1971
1972 let constant = Expr::NumberLiteral(ExprNumberLiteral {
1973 value: Number::Int(Int::ONE),
1974 range: TextRange::default(),
1975 node_index: AtomicNodeIndex::NONE,
1976 });
1977
1978 let type_var_with_bound = TypeParam::TypeVar(TypeParamTypeVar {
1979 range: TextRange::default(),
1980 node_index: AtomicNodeIndex::NONE,
1981 bound: Some(Box::new(constant.clone())),
1982 default: None,
1983 name: Identifier::new("x", TextRange::default()),
1984 });
1985 assert!(
1986 any_over_type_param(&type_var_with_bound, &mut |expr| {
1987 assert_eq!(
1988 *expr, constant,
1989 "the received expression should be the unwrapped bound"
1990 );
1991 true
1992 }),
1993 "if true is returned from `func` it should be respected"
1994 );
1995
1996 let type_var_with_default = TypeParam::TypeVar(TypeParamTypeVar {
1997 range: TextRange::default(),
1998 node_index: AtomicNodeIndex::NONE,
1999 default: Some(Box::new(constant.clone())),
2000 bound: None,
2001 name: Identifier::new("x", TextRange::default()),
2002 });
2003 assert!(
2004 any_over_type_param(&type_var_with_default, &mut |expr| {
2005 assert_eq!(
2006 *expr, constant,
2007 "the received expression should be the unwrapped default"
2008 );
2009 true
2010 }),
2011 "if true is returned from `func` it should be respected"
2012 );
2013 }
2014
2015 #[test]
2016 fn any_over_type_param_type_var_tuple() {
2017 let type_var_tuple = TypeParam::TypeVarTuple(TypeParamTypeVarTuple {
2018 range: TextRange::default(),
2019 node_index: AtomicNodeIndex::NONE,
2020 name: Identifier::new("x", TextRange::default()),
2021 default: None,
2022 });
2023 assert!(
2024 !any_over_type_param(&type_var_tuple, &mut |_expr| true),
2025 "this TypeVarTuple has no expressions to visit"
2026 );
2027
2028 let constant = Expr::NumberLiteral(ExprNumberLiteral {
2029 value: Number::Int(Int::ONE),
2030 range: TextRange::default(),
2031 node_index: AtomicNodeIndex::NONE,
2032 });
2033
2034 let type_var_tuple_with_default = TypeParam::TypeVarTuple(TypeParamTypeVarTuple {
2035 range: TextRange::default(),
2036 node_index: AtomicNodeIndex::NONE,
2037 default: Some(Box::new(constant.clone())),
2038 name: Identifier::new("x", TextRange::default()),
2039 });
2040 assert!(
2041 any_over_type_param(&type_var_tuple_with_default, &mut |expr| {
2042 assert_eq!(
2043 *expr, constant,
2044 "the received expression should be the unwrapped default"
2045 );
2046 true
2047 }),
2048 "if true is returned from `func` it should be respected"
2049 );
2050 }
2051
2052 #[test]
2053 fn any_over_type_param_param_spec() {
2054 let type_param_spec = TypeParam::ParamSpec(TypeParamParamSpec {
2055 range: TextRange::default(),
2056 node_index: AtomicNodeIndex::NONE,
2057 name: Identifier::new("x", TextRange::default()),
2058 default: None,
2059 });
2060 assert!(
2061 !any_over_type_param(&type_param_spec, &mut |_expr| true),
2062 "this ParamSpec has no expressions to visit"
2063 );
2064
2065 let constant = Expr::NumberLiteral(ExprNumberLiteral {
2066 value: Number::Int(Int::ONE),
2067 range: TextRange::default(),
2068 node_index: AtomicNodeIndex::NONE,
2069 });
2070
2071 let param_spec_with_default = TypeParam::TypeVarTuple(TypeParamTypeVarTuple {
2072 range: TextRange::default(),
2073 node_index: AtomicNodeIndex::NONE,
2074 default: Some(Box::new(constant.clone())),
2075 name: Identifier::new("x", TextRange::default()),
2076 });
2077 assert!(
2078 any_over_type_param(¶m_spec_with_default, &mut |expr| {
2079 assert_eq!(
2080 *expr, constant,
2081 "the received expression should be the unwrapped default"
2082 );
2083 true
2084 }),
2085 "if true is returned from `func` it should be respected"
2086 );
2087 }
2088}