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