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

ruff_python_ast/
node.rs
