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

ruff_python_ast/
node.rs
