ruff_python_ast/
node.rs

1use ruff_text_size::Ranged;
2
3use crate::visitor::source_order::SourceOrderVisitor;
4use crate::{
5    self as ast, Alias, AnyNodeRef, AnyParameterRef, ArgOrKeyword, MatchCase, PatternArguments,
6    PatternKeyword,
7};
8
9impl ast::ElifElseClause {
10    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
11    where
12        V: SourceOrderVisitor<'a> + ?Sized,
13    {
14        let ast::ElifElseClause {
15            range: _,
16            node_index: _,
17            test,
18            body,
19        } = self;
20        if let Some(test) = test {
21            visitor.visit_expr(test);
22        }
23        visitor.visit_body(body);
24    }
25}
26
27impl ast::ExprDict {
28    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
29    where
30        V: SourceOrderVisitor<'a> + ?Sized,
31    {
32        let ast::ExprDict {
33            items,
34            range: _,
35            node_index: _,
36        } = self;
37
38        for ast::DictItem { key, value } in items {
39            if let Some(key) = key {
40                visitor.visit_expr(key);
41            }
42            visitor.visit_expr(value);
43        }
44    }
45}
46
47impl ast::ExprBoolOp {
48    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
49    where
50        V: SourceOrderVisitor<'a> + ?Sized,
51    {
52        let ast::ExprBoolOp {
53            op,
54            values,
55            range: _,
56            node_index: _,
57        } = self;
58        match values.as_slice() {
59            [left, rest @ ..] => {
60                visitor.visit_expr(left);
61                visitor.visit_bool_op(op);
62                for expr in rest {
63                    visitor.visit_expr(expr);
64                }
65            }
66            [] => {
67                visitor.visit_bool_op(op);
68            }
69        }
70    }
71}
72
73impl ast::ExprCompare {
74    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
75    where
76        V: SourceOrderVisitor<'a> + ?Sized,
77    {
78        let ast::ExprCompare {
79            left,
80            ops,
81            comparators,
82            range: _,
83            node_index: _,
84        } = self;
85
86        visitor.visit_expr(left);
87
88        for (op, comparator) in ops.iter().zip(comparators) {
89            visitor.visit_cmp_op(op);
90            visitor.visit_expr(comparator);
91        }
92    }
93}
94
95impl ast::InterpolatedStringFormatSpec {
96    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
97    where
98        V: SourceOrderVisitor<'a> + ?Sized,
99    {
100        for element in &self.elements {
101            visitor.visit_interpolated_string_element(element);
102        }
103    }
104}
105
106impl ast::InterpolatedElement {
107    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
108    where
109        V: SourceOrderVisitor<'a> + ?Sized,
110    {
111        let ast::InterpolatedElement {
112            expression,
113            format_spec,
114            ..
115        } = self;
116        visitor.visit_expr(expression);
117
118        if let Some(format_spec) = format_spec {
119            for spec_part in &format_spec.elements {
120                visitor.visit_interpolated_string_element(spec_part);
121            }
122        }
123    }
124}
125
126impl ast::InterpolatedStringLiteralElement {
127    pub(crate) fn visit_source_order<'a, V>(&'a self, _visitor: &mut V)
128    where
129        V: SourceOrderVisitor<'a> + ?Sized,
130    {
131        let ast::InterpolatedStringLiteralElement {
132            range: _,
133            node_index: _,
134            value: _,
135        } = self;
136    }
137}
138
139impl ast::ExprFString {
140    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
141    where
142        V: SourceOrderVisitor<'a> + ?Sized,
143    {
144        let ast::ExprFString {
145            value,
146            range: _,
147            node_index: _,
148        } = self;
149
150        for f_string_part in value {
151            match f_string_part {
152                ast::FStringPart::Literal(string_literal) => {
153                    visitor.visit_string_literal(string_literal);
154                }
155                ast::FStringPart::FString(f_string) => {
156                    visitor.visit_f_string(f_string);
157                }
158            }
159        }
160    }
161}
162
163impl ast::ExprTString {
164    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
165    where
166        V: SourceOrderVisitor<'a> + ?Sized,
167    {
168        let ast::ExprTString {
169            value,
170            range: _,
171            node_index: _,
172        } = self;
173
174        for t_string in value {
175            visitor.visit_t_string(t_string);
176        }
177    }
178}
179
180impl ast::ExprStringLiteral {
181    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
182    where
183        V: SourceOrderVisitor<'a> + ?Sized,
184    {
185        let ast::ExprStringLiteral {
186            value,
187            range: _,
188            node_index: _,
189        } = self;
190
191        for string_literal in value {
192            visitor.visit_string_literal(string_literal);
193        }
194    }
195}
196
197impl ast::ExprBytesLiteral {
198    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
199    where
200        V: SourceOrderVisitor<'a> + ?Sized,
201    {
202        let ast::ExprBytesLiteral {
203            value,
204            range: _,
205            node_index: _,
206        } = self;
207
208        for bytes_literal in value {
209            visitor.visit_bytes_literal(bytes_literal);
210        }
211    }
212}
213
214impl ast::ExceptHandlerExceptHandler {
215    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
216    where
217        V: SourceOrderVisitor<'a> + ?Sized,
218    {
219        let ast::ExceptHandlerExceptHandler {
220            range: _,
221            node_index: _,
222            type_,
223            name,
224            body,
225        } = self;
226        if let Some(expr) = type_ {
227            visitor.visit_expr(expr);
228        }
229
230        if let Some(name) = name {
231            visitor.visit_identifier(name);
232        }
233
234        visitor.visit_body(body);
235    }
236}
237
238impl ast::PatternMatchMapping {
239    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
240    where
241        V: SourceOrderVisitor<'a> + ?Sized,
242    {
243        let ast::PatternMatchMapping {
244            keys,
245            patterns,
246            rest,
247            range: _,
248            node_index: _,
249        } = self;
250
251        let mut rest = rest.as_ref();
252
253        for (key, pattern) in keys.iter().zip(patterns) {
254            if let Some(rest_identifier) = rest
255                && rest_identifier.start() < key.start()
256            {
257                visitor.visit_identifier(rest_identifier);
258                rest = None;
259            }
260            visitor.visit_expr(key);
261            visitor.visit_pattern(pattern);
262        }
263
264        if let Some(rest) = rest {
265            visitor.visit_identifier(rest);
266        }
267    }
268}
269
270impl ast::PatternArguments {
271    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
272    where
273        V: SourceOrderVisitor<'a> + ?Sized,
274    {
275        let PatternArguments {
276            range: _,
277            node_index: _,
278            patterns,
279            keywords,
280        } = self;
281
282        for pattern in patterns {
283            visitor.visit_pattern(pattern);
284        }
285
286        for keyword in keywords {
287            visitor.visit_pattern_keyword(keyword);
288        }
289    }
290}
291
292impl ast::PatternKeyword {
293    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
294    where
295        V: SourceOrderVisitor<'a> + ?Sized,
296    {
297        let PatternKeyword {
298            range: _,
299            node_index: _,
300            attr,
301            pattern,
302        } = self;
303
304        visitor.visit_identifier(attr);
305        visitor.visit_pattern(pattern);
306    }
307}
308
309impl ast::Comprehension {
310    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
311    where
312        V: SourceOrderVisitor<'a> + ?Sized,
313    {
314        let ast::Comprehension {
315            range: _,
316            node_index: _,
317            target,
318            iter,
319            ifs,
320            is_async: _,
321        } = self;
322        visitor.visit_expr(target);
323        visitor.visit_expr(iter);
324
325        for expr in ifs {
326            visitor.visit_expr(expr);
327        }
328    }
329}
330
331impl ast::Arguments {
332    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
333    where
334        V: SourceOrderVisitor<'a> + ?Sized,
335    {
336        for arg_or_keyword in self.arguments_source_order() {
337            match arg_or_keyword {
338                ArgOrKeyword::Arg(arg) => visitor.visit_expr(arg),
339                ArgOrKeyword::Keyword(keyword) => visitor.visit_keyword(keyword),
340            }
341        }
342    }
343}
344
345impl ast::Parameters {
346    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
347    where
348        V: SourceOrderVisitor<'a> + ?Sized,
349    {
350        for parameter in self {
351            match parameter {
352                AnyParameterRef::NonVariadic(parameter_with_default) => {
353                    visitor.visit_parameter_with_default(parameter_with_default);
354                }
355                AnyParameterRef::Variadic(parameter) => visitor.visit_parameter(parameter),
356            }
357        }
358    }
359}
360
361impl ast::Parameter {
362    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
363    where
364        V: SourceOrderVisitor<'a> + ?Sized,
365    {
366        let ast::Parameter {
367            range: _,
368            node_index: _,
369            name,
370            annotation,
371        } = self;
372
373        visitor.visit_identifier(name);
374        if let Some(expr) = annotation {
375            visitor.visit_annotation(expr);
376        }
377    }
378}
379
380impl ast::ParameterWithDefault {
381    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
382    where
383        V: SourceOrderVisitor<'a> + ?Sized,
384    {
385        let ast::ParameterWithDefault {
386            range: _,
387            node_index: _,
388            parameter,
389            default,
390        } = self;
391        visitor.visit_parameter(parameter);
392        if let Some(expr) = default {
393            visitor.visit_expr(expr);
394        }
395    }
396}
397
398impl ast::Keyword {
399    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
400    where
401        V: SourceOrderVisitor<'a> + ?Sized,
402    {
403        let ast::Keyword {
404            range: _,
405            node_index: _,
406            arg,
407            value,
408        } = self;
409
410        if let Some(arg) = arg {
411            visitor.visit_identifier(arg);
412        }
413        visitor.visit_expr(value);
414    }
415}
416
417impl Alias {
418    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
419    where
420        V: SourceOrderVisitor<'a> + ?Sized,
421    {
422        let ast::Alias {
423            range: _,
424            node_index: _,
425            name,
426            asname,
427        } = self;
428
429        visitor.visit_identifier(name);
430        if let Some(asname) = asname {
431            visitor.visit_identifier(asname);
432        }
433    }
434}
435
436impl ast::WithItem {
437    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
438    where
439        V: SourceOrderVisitor<'a> + ?Sized,
440    {
441        let ast::WithItem {
442            range: _,
443            node_index: _,
444            context_expr,
445            optional_vars,
446        } = self;
447
448        visitor.visit_expr(context_expr);
449
450        if let Some(expr) = optional_vars {
451            visitor.visit_expr(expr);
452        }
453    }
454}
455
456impl ast::MatchCase {
457    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
458    where
459        V: SourceOrderVisitor<'a> + ?Sized,
460    {
461        let ast::MatchCase {
462            range: _,
463            node_index: _,
464            pattern,
465            guard,
466            body,
467        } = self;
468
469        visitor.visit_pattern(pattern);
470        if let Some(expr) = guard {
471            visitor.visit_expr(expr);
472        }
473        visitor.visit_body(body);
474    }
475}
476
477impl ast::Decorator {
478    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
479    where
480        V: SourceOrderVisitor<'a> + ?Sized,
481    {
482        let ast::Decorator {
483            range: _,
484            node_index: _,
485            expression,
486        } = self;
487
488        visitor.visit_expr(expression);
489    }
490}
491
492impl ast::TypeParams {
493    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
494    where
495        V: SourceOrderVisitor<'a> + ?Sized,
496    {
497        let ast::TypeParams {
498            range: _,
499            node_index: _,
500            type_params,
501        } = self;
502
503        for type_param in type_params {
504            visitor.visit_type_param(type_param);
505        }
506    }
507}
508
509impl ast::FString {
510    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
511    where
512        V: SourceOrderVisitor<'a> + ?Sized,
513    {
514        let ast::FString {
515            elements,
516            range: _,
517            node_index: _,
518            flags: _,
519        } = self;
520
521        for fstring_element in elements {
522            visitor.visit_interpolated_string_element(fstring_element);
523        }
524    }
525}
526
527impl ast::TString {
528    pub(crate) fn visit_source_order<'a, V>(&'a self, visitor: &mut V)
529    where
530        V: SourceOrderVisitor<'a> + ?Sized,
531    {
532        let ast::TString {
533            elements,
534            range: _,
535            node_index: _,
536            flags: _,
537        } = self;
538
539        for tstring_element in elements {
540            visitor.visit_interpolated_string_element(tstring_element);
541        }
542    }
543}
544
545impl ast::StringLiteral {
546    #[inline]
547    pub(crate) fn visit_source_order<'a, V>(&'a self, _visitor: &mut V)
548    where
549        V: SourceOrderVisitor<'a> + ?Sized,
550    {
551        let ast::StringLiteral {
552            range: _,
553            node_index: _,
554            value: _,
555            flags: _,
556        } = self;
557    }
558}
559
560impl ast::BytesLiteral {
561    #[inline]
562    pub(crate) fn visit_source_order<'a, V>(&'a self, _visitor: &mut V)
563    where
564        V: SourceOrderVisitor<'a> + ?Sized,
565    {
566        let ast::BytesLiteral {
567            range: _,
568            node_index: _,
569            value: _,
570            flags: _,
571        } = self;
572    }
573}
574
575impl ast::Identifier {
576    #[inline]
577    pub(crate) fn visit_source_order<'a, V>(&'a self, _visitor: &mut V)
578    where
579        V: SourceOrderVisitor<'a> + ?Sized,
580    {
581        let ast::Identifier {
582            range: _,
583            node_index: _,
584            id: _,
585        } = self;
586    }
587}
588
589impl<'a> AnyNodeRef<'a> {
590    /// Compares two any node refs by their pointers (referential equality).
591    pub fn ptr_eq(self, other: AnyNodeRef) -> bool {
592        self.as_ptr().eq(&other.as_ptr()) && self.kind() == other.kind()
593    }
594
595    /// In our AST, only some alternative branches are represented as a node. This has historical
596    /// reasons, e.g. we added a node for elif/else in if statements which was not originally
597    /// present in the parser.
598    pub const fn is_alternative_branch_with_node(self) -> bool {
599        matches!(
600            self,
601            AnyNodeRef::ExceptHandlerExceptHandler(_) | AnyNodeRef::ElifElseClause(_)
602        )
603    }
604
605    /// The last child of the last branch, if the node has multiple branches.
606    pub fn last_child_in_body(&self) -> Option<AnyNodeRef<'a>> {
607        let body =
608            match self {
609                AnyNodeRef::StmtFunctionDef(ast::StmtFunctionDef { body, .. })
610                | AnyNodeRef::StmtClassDef(ast::StmtClassDef { body, .. })
611                | AnyNodeRef::StmtWith(ast::StmtWith { body, .. })
612                | AnyNodeRef::MatchCase(MatchCase { body, .. })
613                | AnyNodeRef::ExceptHandlerExceptHandler(ast::ExceptHandlerExceptHandler {
614                    body,
615                    ..
616                })
617                | AnyNodeRef::ElifElseClause(ast::ElifElseClause { body, .. }) => body,
618                AnyNodeRef::StmtIf(ast::StmtIf {
619                    body,
620                    elif_else_clauses,
621                    ..
622                }) => elif_else_clauses.last().map_or(body, |clause| &clause.body),
623
624                AnyNodeRef::StmtFor(ast::StmtFor { body, orelse, .. })
625                | AnyNodeRef::StmtWhile(ast::StmtWhile { body, orelse, .. }) => {
626                    if orelse.is_empty() { body } else { orelse }
627                }
628
629                AnyNodeRef::StmtMatch(ast::StmtMatch { cases, .. }) => {
630                    return cases.last().map(AnyNodeRef::from);
631                }
632
633                AnyNodeRef::StmtTry(ast::StmtTry {
634                    body,
635                    handlers,
636                    orelse,
637                    finalbody,
638                    ..
639                }) => {
640                    if finalbody.is_empty() {
641                        if orelse.is_empty() {
642                            if handlers.is_empty() {
643                                body
644                            } else {
645                                return handlers.last().map(AnyNodeRef::from);
646                            }
647                        } else {
648                            orelse
649                        }
650                    } else {
651                        finalbody
652                    }
653                }
654
655                // Not a node that contains an indented child node.
656                _ => return None,
657            };
658
659        body.last().map(AnyNodeRef::from)
660    }
661
662    /// Check if the given statement is the first statement after the colon of a branch, be it in if
663    /// statements, for statements, after each part of a try-except-else-finally or function/class
664    /// definitions.
665    ///
666    ///
667    /// ```python
668    /// if True:    <- has body
669    ///     a       <- first statement
670    ///     b
671    /// elif b:     <- has body
672    ///     c       <- first statement
673    ///     d
674    /// else:       <- has body
675    ///     e       <- first statement
676    ///     f
677    ///
678    /// class:      <- has body
679    ///     a: int  <- first statement
680    ///     b: int
681    ///
682    /// ```
683    ///
684    /// For nodes with multiple bodies, we check all bodies that don't have their own node. For
685    /// try-except-else-finally, each except branch has it's own node, so for the `StmtTry`, we check
686    /// the `try:`, `else:` and `finally:`, bodies, while `ExceptHandlerExceptHandler` has it's own
687    /// check. For for-else and while-else, we check both branches for the whole statement.
688    ///
689    /// ```python
690    /// try:        <- has body (a)
691    ///     6/8     <- first statement (a)
692    ///     1/0
693    /// except:     <- has body (b)
694    ///     a       <- first statement (b)
695    ///     b
696    /// else:
697    ///     c       <- first statement (a)
698    ///     d
699    /// finally:
700    ///     e       <- first statement (a)
701    ///     f
702    /// ```
703    pub fn is_first_statement_in_body(&self, body: AnyNodeRef) -> bool {
704        match body {
705            AnyNodeRef::StmtFor(ast::StmtFor { body, orelse, .. })
706            | AnyNodeRef::StmtWhile(ast::StmtWhile { body, orelse, .. }) => {
707                are_same_optional(*self, body.first()) || are_same_optional(*self, orelse.first())
708            }
709
710            AnyNodeRef::StmtTry(ast::StmtTry {
711                body,
712                orelse,
713                finalbody,
714                ..
715            }) => {
716                are_same_optional(*self, body.first())
717                    || are_same_optional(*self, orelse.first())
718                    || are_same_optional(*self, finalbody.first())
719            }
720
721            AnyNodeRef::StmtIf(ast::StmtIf { body, .. })
722            | AnyNodeRef::ElifElseClause(ast::ElifElseClause { body, .. })
723            | AnyNodeRef::StmtWith(ast::StmtWith { body, .. })
724            | AnyNodeRef::ExceptHandlerExceptHandler(ast::ExceptHandlerExceptHandler {
725                body,
726                ..
727            })
728            | AnyNodeRef::MatchCase(MatchCase { body, .. })
729            | AnyNodeRef::StmtFunctionDef(ast::StmtFunctionDef { body, .. })
730            | AnyNodeRef::StmtClassDef(ast::StmtClassDef { body, .. }) => {
731                are_same_optional(*self, body.first())
732            }
733
734            AnyNodeRef::StmtMatch(ast::StmtMatch { cases, .. }) => {
735                are_same_optional(*self, cases.first())
736            }
737
738            _ => false,
739        }
740    }
741
742    /// Returns `true` if `statement` is the first statement in an alternate `body` (e.g. the else of an if statement)
743    pub fn is_first_statement_in_alternate_body(&self, body: AnyNodeRef) -> bool {
744        match body {
745            AnyNodeRef::StmtFor(ast::StmtFor { orelse, .. })
746            | AnyNodeRef::StmtWhile(ast::StmtWhile { orelse, .. }) => {
747                are_same_optional(*self, orelse.first())
748            }
749
750            AnyNodeRef::StmtTry(ast::StmtTry {
751                handlers,
752                orelse,
753                finalbody,
754                ..
755            }) => {
756                are_same_optional(*self, handlers.first())
757                    || are_same_optional(*self, orelse.first())
758                    || are_same_optional(*self, finalbody.first())
759            }
760
761            AnyNodeRef::StmtIf(ast::StmtIf {
762                elif_else_clauses, ..
763            }) => are_same_optional(*self, elif_else_clauses.first()),
764            _ => false,
765        }
766    }
767}
768
769/// Returns `true` if `right` is `Some` and `left` and `right` are referentially equal.
770fn are_same_optional<'a, T>(left: AnyNodeRef, right: Option<T>) -> bool
771where
772    T: Into<AnyNodeRef<'a>>,
773{
774    right.is_some_and(|right| left.ptr_eq(right.into()))
775}
ruff_python_ast/node.rs

ruff_python_ast/
node.rs
