Skip to main content

ruff_python_ast/
nodes.rs

1#![allow(clippy::derive_partial_eq_without_eq)]
2
3use crate::AtomicNodeIndex;
4use crate::generated::{
5    ExprBytesLiteral, ExprCall, ExprCompare, ExprDict, ExprFString, ExprList, ExprName, ExprSet,
6    ExprStringLiteral, ExprTString, ExprTuple, PatternMatchAs, PatternMatchOr, StmtClassDef,
7};
8use std::borrow::Cow;
9use std::fmt;
10use std::fmt::Debug;
11use std::iter::{FusedIterator, Once};
12use std::ops::{Deref, DerefMut};
13use std::slice::{Iter, IterMut};
14use std::sync::OnceLock;
15
16use bitflags::bitflags;
17use thin_vec::ThinVec;
18
19use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
20
21use crate::str_prefix::{
22    AnyStringPrefix, ByteStringPrefix, FStringPrefix, StringLiteralPrefix, TStringPrefix,
23};
24use crate::{
25    Expr, ExprRef, InterpolatedStringElement, LiteralExpressionRef, OperatorPrecedence, Pattern,
26    Stmt, TypeParam, int,
27    name::Name,
28    str::{Quote, TripleQuotes},
29};
30
31#[derive(Clone, Debug, PartialEq)]
32#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
33pub enum ConstantValue {
34    None,
35    Boolean(bool),
36    Str(Box<str>),
37    Bytes(Box<[u8]>),
38    Integer(Box<str>),
39    Tuple(Vec<Self>),
40    Frozenset(Vec<Self>),
41    Float(f64),
42    Complex { real: f64, imag: f64 },
43    Ellipsis,
44}
45
46impl StmtClassDef {
47    /// Return an iterator over the bases of the class.
48    pub fn bases(&self) -> &[Expr] {
49        match &self.arguments {
50            Some(arguments) => &arguments.args,
51            None => &[],
52        }
53    }
54
55    /// Return an iterator over the metaclass keywords of the class.
56    pub fn keywords(&self) -> &[Keyword] {
57        match &self.arguments {
58            Some(arguments) => &arguments.keywords,
59            None => &[],
60        }
61    }
62}
63
64#[derive(Clone, Debug, PartialEq)]
65#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
66pub struct ElifElseClause {
67    pub range: TextRange,
68    pub node_index: AtomicNodeIndex,
69    pub test: Option<Expr>,
70    pub body: Suite,
71    pub runtime_body: Option<Vec<Option<Stmt>>>,
72    pub runtime_orelse: Option<Vec<Option<Stmt>>>,
73}
74
75impl Expr {
76    /// Returns `true` if the expression is a literal expression.
77    ///
78    /// A literal expression is either a string literal, bytes literal,
79    /// integer, float, complex number, boolean, `None`, or ellipsis (`...`).
80    pub fn is_literal_expr(&self) -> bool {
81        matches!(
82            self,
83            Expr::StringLiteral(_)
84                | Expr::BytesLiteral(_)
85                | Expr::NumberLiteral(_)
86                | Expr::Constant(_)
87                | Expr::BooleanLiteral(_)
88                | Expr::NoneLiteral(_)
89                | Expr::EllipsisLiteral(_)
90        )
91    }
92
93    /// Returns [`LiteralExpressionRef`] if the expression is a literal expression.
94    pub fn as_literal_expr(&self) -> Option<LiteralExpressionRef<'_>> {
95        match self {
96            Expr::StringLiteral(expr) => Some(LiteralExpressionRef::StringLiteral(expr)),
97            Expr::BytesLiteral(expr) => Some(LiteralExpressionRef::BytesLiteral(expr)),
98            Expr::NumberLiteral(expr) => Some(LiteralExpressionRef::NumberLiteral(expr)),
99            Expr::BooleanLiteral(expr) => Some(LiteralExpressionRef::BooleanLiteral(expr)),
100            Expr::NoneLiteral(expr) => Some(LiteralExpressionRef::NoneLiteral(expr)),
101            Expr::EllipsisLiteral(expr) => Some(LiteralExpressionRef::EllipsisLiteral(expr)),
102            _ => None,
103        }
104    }
105
106    /// Return the value expression after peeling off any nested named expressions.
107    ///
108    /// For example, this returns the `x` expression for both `x` and `(y := x)`.
109    pub fn expression_value(&self) -> &Self {
110        let mut expr = self;
111        while let Expr::Named(named) = expr {
112            expr = &named.value;
113        }
114        expr
115    }
116
117    /// Return the [`OperatorPrecedence`] of this expression
118    pub fn precedence(&self) -> OperatorPrecedence {
119        OperatorPrecedence::from(self)
120    }
121}
122
123impl ExprCompare {
124    /// Returns the initial operand (`a` in `a < b <= c`).
125    ///
126    /// # Panics
127    ///
128    /// Panics if there are no operands. The parser always produces at least two operands,
129    /// synthesizing an invalid name expression for a missing operand during error recovery.
130    pub fn first_operand(&self) -> &Expr {
131        self.operands
132            .first()
133            .expect("A comparison has a left operand")
134    }
135
136    /// Returns the right operand of the first comparison (`b` in `a < b <= c`).
137    ///
138    /// # Panics
139    ///
140    /// Panics if there are fewer than two operands. See [`Self::first_operand`].
141    pub fn second_operand(&self) -> &Expr {
142        self.operands
143            .get(1)
144            .expect("A comparison has a right operand")
145    }
146
147    /// Returns the first comparison operator (`<` in `a < b <= c`).
148    ///
149    /// # Panics
150    ///
151    /// Panics if there are no operators. The parser always produces at least one operator.
152    pub fn first_operator(&self) -> CmpOp {
153        *self.ops.first().expect("A comparison has an operator")
154    }
155
156    /// Returns all operands after the initial operand (`[b, c]` in `a < b <= c`).
157    ///
158    /// # Panics
159    ///
160    /// Panics if there are no operands. See [`Self::first_operand`].
161    pub fn comparators(&self) -> &[Expr] {
162        &self.operands[1..]
163    }
164
165    /// Returns `(left, operator, right)` if there is exactly one operator and two operands.
166    pub fn as_single(&self) -> Option<(&Expr, &CmpOp, &Expr)> {
167        match (&*self.operands, &*self.ops) {
168            ([left, right], [op]) => Some((left, op, right)),
169            _ => None,
170        }
171    }
172
173    /// Iterates over each comparison as `(left, operator, right)`.
174    ///
175    /// For `a < b <= c`, yields `(a, <, b)` followed by `(b, <=, c)`.
176    pub fn iter(
177        &self,
178    ) -> impl DoubleEndedIterator<Item = (&Expr, &CmpOp, &Expr)> + ExactSizeIterator {
179        self.operands
180            .iter()
181            .zip(self.ops.iter())
182            .zip(self.operands.iter().skip(1))
183            .map(|((left, op), right)| (left, op, right))
184    }
185}
186
187impl ExprRef<'_> {
188    /// See [`Expr::is_literal_expr`].
189    pub fn is_literal_expr(&self) -> bool {
190        matches!(
191            self,
192            ExprRef::StringLiteral(_)
193                | ExprRef::BytesLiteral(_)
194                | ExprRef::NumberLiteral(_)
195                | ExprRef::BooleanLiteral(_)
196                | ExprRef::NoneLiteral(_)
197                | ExprRef::EllipsisLiteral(_)
198        )
199    }
200
201    pub fn precedence(&self) -> OperatorPrecedence {
202        OperatorPrecedence::from(*self)
203    }
204}
205
206/// Represents an item in a [dictionary literal display][1].
207///
208/// Consider the following Python dictionary literal:
209/// ```python
210/// {key1: value1, **other_dictionary}
211/// ```
212///
213/// In our AST, this would be represented using an `ExprDict` node containing
214/// two `DictItem` nodes inside it:
215/// ```ignore
216/// [
217///     DictItem {
218///         key: Some(Expr::Name(ExprName { id: "key1" })),
219///         value: Expr::Name(ExprName { id: "value1" }),
220///     },
221///     DictItem {
222///         key: None,
223///         value: Expr::Name(ExprName { id: "other_dictionary" }),
224///     }
225/// ]
226/// ```
227///
228/// [1]: https://docs.python.org/3/reference/expressions.html#displays-for-lists-sets-and-dictionaries
229#[derive(Debug, Clone, PartialEq)]
230#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
231pub struct DictItem {
232    pub key: Option<Expr>,
233    pub value: Expr,
234}
235
236impl DictItem {
237    fn key(&self) -> Option<&Expr> {
238        self.key.as_ref()
239    }
240
241    fn value(&self) -> &Expr {
242        &self.value
243    }
244}
245
246impl Ranged for DictItem {
247    fn range(&self) -> TextRange {
248        TextRange::new(
249            self.key.as_ref().map_or(self.value.start(), Ranged::start),
250            self.value.end(),
251        )
252    }
253}
254
255impl ExprDict {
256    /// Returns an `Iterator` over the AST nodes representing the
257    /// dictionary's keys.
258    pub fn iter_keys(&self) -> DictKeyIterator<'_> {
259        DictKeyIterator::new(&self.items)
260    }
261
262    /// Returns an `Iterator` over the AST nodes representing the
263    /// dictionary's values.
264    pub fn iter_values(&self) -> DictValueIterator<'_> {
265        DictValueIterator::new(&self.items)
266    }
267
268    /// Returns the AST node representing the *n*th key of this
269    /// dictionary.
270    ///
271    /// Panics: If the index `n` is out of bounds.
272    pub fn key(&self, n: usize) -> Option<&Expr> {
273        self.items[n].key()
274    }
275
276    /// Returns the AST node representing the *n*th value of this
277    /// dictionary.
278    ///
279    /// Panics: If the index `n` is out of bounds.
280    pub fn value(&self, n: usize) -> &Expr {
281        self.items[n].value()
282    }
283
284    pub fn iter(&self) -> std::slice::Iter<'_, DictItem> {
285        self.items.iter()
286    }
287
288    pub fn len(&self) -> usize {
289        self.items.len()
290    }
291
292    pub fn is_empty(&self) -> bool {
293        self.items.is_empty()
294    }
295}
296
297impl<'a> IntoIterator for &'a ExprDict {
298    type IntoIter = std::slice::Iter<'a, DictItem>;
299    type Item = &'a DictItem;
300
301    fn into_iter(self) -> Self::IntoIter {
302        self.iter()
303    }
304}
305
306#[derive(Debug, Clone)]
307pub struct DictKeyIterator<'a> {
308    items: Iter<'a, DictItem>,
309}
310
311impl<'a> DictKeyIterator<'a> {
312    fn new(items: &'a [DictItem]) -> Self {
313        Self {
314            items: items.iter(),
315        }
316    }
317
318    pub fn is_empty(&self) -> bool {
319        self.len() == 0
320    }
321}
322
323impl<'a> Iterator for DictKeyIterator<'a> {
324    type Item = Option<&'a Expr>;
325
326    fn next(&mut self) -> Option<Self::Item> {
327        self.items.next().map(DictItem::key)
328    }
329
330    fn last(mut self) -> Option<Self::Item> {
331        self.next_back()
332    }
333
334    fn size_hint(&self) -> (usize, Option<usize>) {
335        self.items.size_hint()
336    }
337}
338
339impl DoubleEndedIterator for DictKeyIterator<'_> {
340    fn next_back(&mut self) -> Option<Self::Item> {
341        self.items.next_back().map(DictItem::key)
342    }
343}
344
345impl FusedIterator for DictKeyIterator<'_> {}
346impl ExactSizeIterator for DictKeyIterator<'_> {}
347
348#[derive(Debug, Clone)]
349pub struct DictValueIterator<'a> {
350    items: Iter<'a, DictItem>,
351}
352
353impl<'a> DictValueIterator<'a> {
354    fn new(items: &'a [DictItem]) -> Self {
355        Self {
356            items: items.iter(),
357        }
358    }
359
360    pub fn is_empty(&self) -> bool {
361        self.len() == 0
362    }
363}
364
365impl<'a> Iterator for DictValueIterator<'a> {
366    type Item = &'a Expr;
367
368    fn next(&mut self) -> Option<Self::Item> {
369        self.items.next().map(DictItem::value)
370    }
371
372    fn last(mut self) -> Option<Self::Item> {
373        self.next_back()
374    }
375
376    fn size_hint(&self) -> (usize, Option<usize>) {
377        self.items.size_hint()
378    }
379}
380
381impl DoubleEndedIterator for DictValueIterator<'_> {
382    fn next_back(&mut self) -> Option<Self::Item> {
383        self.items.next_back().map(DictItem::value)
384    }
385}
386
387impl FusedIterator for DictValueIterator<'_> {}
388impl ExactSizeIterator for DictValueIterator<'_> {}
389
390impl ExprSet {
391    pub fn iter(&self) -> std::slice::Iter<'_, Expr> {
392        self.elts.iter()
393    }
394
395    pub fn len(&self) -> usize {
396        self.elts.len()
397    }
398
399    pub fn is_empty(&self) -> bool {
400        self.elts.is_empty()
401    }
402}
403
404impl<'a> IntoIterator for &'a ExprSet {
405    type IntoIter = std::slice::Iter<'a, Expr>;
406    type Item = &'a Expr;
407
408    fn into_iter(self) -> Self::IntoIter {
409        self.iter()
410    }
411}
412
413#[derive(Clone, Debug, PartialEq)]
414#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
415pub struct InterpolatedStringFormatSpec {
416    pub range: TextRange,
417    pub node_index: AtomicNodeIndex,
418    pub elements: InterpolatedStringElements,
419}
420
421/// See also [FormattedValue](https://docs.python.org/3/library/ast.html#ast.FormattedValue)
422#[derive(Clone, Debug, PartialEq)]
423#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
424pub struct InterpolatedElement {
425    pub range: TextRange,
426    pub node_index: AtomicNodeIndex,
427    pub expression: Box<Expr>,
428    pub debug_text: Option<DebugText>,
429    pub conversion: ConversionFlag,
430    pub format_spec: Option<Box<InterpolatedStringFormatSpec>>,
431    pub runtime_str: Option<crate::ConstantValue>,
432    pub runtime_interpolation_format_spec: Option<Box<Expr>>,
433    pub runtime_formatted_value_format_spec: Option<Box<Expr>>,
434}
435
436/// An `FStringLiteralElement` with an empty `value` is an invalid f-string element.
437#[derive(Clone, Debug, PartialEq)]
438#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
439pub struct InterpolatedStringLiteralElement {
440    pub range: TextRange,
441    pub node_index: AtomicNodeIndex,
442    pub value: Box<str>,
443}
444
445impl InterpolatedStringLiteralElement {
446    pub fn is_valid(&self) -> bool {
447        !self.value.is_empty()
448    }
449}
450
451impl Deref for InterpolatedStringLiteralElement {
452    type Target = str;
453
454    fn deref(&self) -> &Self::Target {
455        &self.value
456    }
457}
458
459/// Transforms a value prior to formatting it.
460#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, is_macro::Is)]
461#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
462#[repr(i8)]
463#[expect(clippy::cast_possible_wrap)]
464pub enum ConversionFlag {
465    /// No conversion
466    None = -1, // CPython uses -1
467    /// Converts by calling `str(<value>)`.
468    Str = b's' as i8,
469    /// Converts by calling `ascii(<value>)`.
470    Ascii = b'a' as i8,
471    /// Converts by calling `repr(<value>)`.
472    Repr = b'r' as i8,
473}
474
475impl ConversionFlag {
476    pub fn to_byte(&self) -> Option<u8> {
477        match self {
478            Self::None => None,
479            flag => Some(*flag as u8),
480        }
481    }
482    pub fn to_char(&self) -> Option<char> {
483        Some(self.to_byte()? as char)
484    }
485}
486
487/// The debug text of a self-documenting f-string expression (e.g., `f"{x=}"`).
488///
489/// Stores the concatenation of leading text, expression source, and trailing text as a single
490/// [`CompactString`], with byte offsets to split them. The offsets are needed because the leading
491/// and trailing portions can contain non-whitespace characters (grouping parentheses, comments in
492/// triple-quoted f-strings) that cannot be distinguished from expression content by scanning.
493///
494/// [`CompactString`]: compact_str::CompactString
495#[derive(Clone, PartialEq, Eq, Hash)]
496#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
497pub struct DebugText {
498    /// The full text between the `{` and the conversion / `format_spec` / `}`.
499    text: compact_str::CompactString,
500    /// Byte offset where the expression source begins.
501    expression_start: u32,
502    /// Byte offset where the expression source ends.
503    expression_end: u32,
504}
505
506impl std::fmt::Debug for DebugText {
507    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
508        f.debug_struct("DebugText")
509            .field("leading", &self.leading())
510            .field("expression", &self.expression())
511            .field("trailing", &self.trailing())
512            .finish()
513    }
514}
515
516impl DebugText {
517    pub fn new(leading: &str, expression: &str, trailing: &str) -> Self {
518        let expression_start = leading.text_len().to_u32();
519        let expression_end = expression_start + expression.text_len().to_u32();
520        let mut buf = compact_str::CompactString::with_capacity(
521            leading.len() + expression.len() + trailing.len(),
522        );
523        buf.push_str(leading);
524        buf.push_str(expression);
525        buf.push_str(trailing);
526        Self {
527            text: buf,
528            expression_start,
529            expression_end,
530        }
531    }
532
533    /// The full debug text between the `{` and the conversion / `format_spec` / `}`.
534    pub fn as_str(&self) -> &str {
535        &self.text
536    }
537
538    /// The text between the `{` and the expression node.
539    pub fn leading(&self) -> &str {
540        &self.text[..self.expression_start as usize]
541    }
542
543    /// The source text of the expression (e.g., `0x0` in `f"{0x0=}"`).
544    pub fn expression(&self) -> &str {
545        &self.text[self.expression_start as usize..self.expression_end as usize]
546    }
547
548    /// The text between the expression and the conversion, the `format_spec`, or the `}`.
549    pub fn trailing(&self) -> &str {
550        &self.text[self.expression_end as usize..]
551    }
552}
553
554impl ExprFString {
555    /// Returns the single [`FString`] if the f-string isn't implicitly concatenated, [`None`]
556    /// otherwise.
557    pub const fn as_single_part_fstring(&self) -> Option<&FString> {
558        match &self.value.inner {
559            FStringValueInner::Single(fstring) => Some(fstring),
560            FStringValueInner::Concatenated(_) => None,
561        }
562    }
563}
564
565/// The value representing an [`ExprFString`].
566#[derive(Clone, Debug, PartialEq)]
567#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
568pub struct FStringValue {
569    inner: FStringValueInner,
570}
571
572impl FStringValue {
573    /// Creates a new f-string literal with a single [`FString`] part.
574    pub fn single(value: FString) -> Self {
575        Self {
576            inner: FStringValueInner::Single(value),
577        }
578    }
579
580    /// Creates a new f-string with the given values that represents an implicitly
581    /// concatenated f-string.
582    ///
583    /// # Panics
584    ///
585    /// Panics if `values` has less than 2 elements.
586    /// Use [`FStringValue::single`] instead.
587    pub fn concatenated(values: Vec<FStringPart>) -> Self {
588        assert!(
589            values.len() > 1,
590            "Use `FStringValue::single` to create single-part f-strings"
591        );
592        Self {
593            inner: FStringValueInner::Concatenated(values),
594        }
595    }
596
597    /// Returns `true` if the f-string is implicitly concatenated, `false` otherwise.
598    pub fn is_implicit_concatenated(&self) -> bool {
599        matches!(self.inner, FStringValueInner::Concatenated(_))
600    }
601
602    /// Returns an iterator over the parts of this f-string, in source order.
603    pub fn iter(&self) -> FStringParts<'_> {
604        match &self.inner {
605            FStringValueInner::Single(fstring) => FStringParts::Single(std::iter::once(fstring)),
606            FStringValueInner::Concatenated(parts) => FStringParts::Concatenated(parts.iter()),
607        }
608    }
609
610    /// Returns a mutable iterator over the parts of this f-string, in source order.
611    pub fn iter_mut(&mut self) -> FStringPartsMut<'_> {
612        match &mut self.inner {
613            FStringValueInner::Single(fstring) => FStringPartsMut::Single(std::iter::once(fstring)),
614            FStringValueInner::Concatenated(parts) => {
615                FStringPartsMut::Concatenated(parts.iter_mut())
616            }
617        }
618    }
619
620    /// Returns an iterator over the [`StringLiteral`] parts contained in this value.
621    ///
622    /// Note that this doesn't recurse into the f-string parts. For example,
623    ///
624    /// ```python
625    /// "foo" f"bar {x}" "baz" f"qux"
626    /// ```
627    ///
628    /// Here, the string literal parts returned would be `"foo"` and `"baz"`.
629    pub fn literals(&self) -> impl Iterator<Item = &StringLiteral> {
630        self.iter().filter_map(|part| match part {
631            FStringPartRef::Literal(literal) => Some(literal),
632            FStringPartRef::FString(_) => None,
633        })
634    }
635
636    /// Returns an iterator over the [`FString`] parts contained in this value.
637    ///
638    /// Note that this doesn't recurse into the f-string parts. For example,
639    ///
640    /// ```python
641    /// "foo" f"bar {x}" "baz" f"qux"
642    /// ```
643    ///
644    /// Here, the f-string parts returned would be `f"bar {x}"` and `f"qux"`.
645    pub fn f_strings(&self) -> impl Iterator<Item = &FString> {
646        self.iter().filter_map(|part| match part {
647            FStringPartRef::FString(fstring) => Some(fstring),
648            FStringPartRef::Literal(_) => None,
649        })
650    }
651
652    /// Returns an iterator over all the [`InterpolatedStringElement`] contained in this value.
653    ///
654    /// An f-string element is what makes up an [`FString`] i.e., it is either a
655    /// string literal or an expression. In the following example,
656    ///
657    /// ```python
658    /// "foo" f"bar {x}" "baz" f"qux"
659    /// ```
660    ///
661    /// The f-string elements returned would be string literal (`"bar "`),
662    /// expression (`x`) and string literal (`"qux"`).
663    pub fn elements(&self) -> impl Iterator<Item = &InterpolatedStringElement> {
664        self.f_strings().flat_map(|fstring| fstring.elements.iter())
665    }
666
667    /// Returns `true` if the node represents an empty f-string literal.
668    ///
669    /// Note that a [`FStringValue`] node will always have >= 1 [`FStringPart`]s inside it.
670    /// This method checks whether the value of the concatenated parts is equal to the empty
671    /// f-string, not whether the f-string has 0 parts inside it.
672    pub fn is_empty_literal(&self) -> bool {
673        match &self.inner {
674            FStringValueInner::Single(fstring) => fstring.elements.is_empty(),
675            FStringValueInner::Concatenated(fstring_parts) => {
676                fstring_parts.iter().all(FStringPart::is_empty_literal)
677            }
678        }
679    }
680}
681
682impl<'a> IntoIterator for &'a FStringValue {
683    type Item = FStringPartRef<'a>;
684    type IntoIter = FStringParts<'a>;
685
686    fn into_iter(self) -> Self::IntoIter {
687        self.iter()
688    }
689}
690
691impl<'a> IntoIterator for &'a mut FStringValue {
692    type Item = FStringPartMut<'a>;
693    type IntoIter = FStringPartsMut<'a>;
694    fn into_iter(self) -> Self::IntoIter {
695        self.iter_mut()
696    }
697}
698
699/// An internal representation of [`FStringValue`].
700#[derive(Clone, Debug, PartialEq)]
701#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
702enum FStringValueInner {
703    /// A single f-string i.e., `f"foo"`.
704    Single(FString),
705
706    /// An implicitly concatenated f-string i.e., `"foo" f"bar {x}"`.
707    Concatenated(Vec<FStringPart>),
708}
709
710/// An f-string part which is either a string literal or an f-string.
711#[derive(Clone, Debug, PartialEq, is_macro::Is)]
712#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
713pub enum FStringPart {
714    Literal(StringLiteral),
715    FString(FString),
716}
717
718impl FStringPart {
719    pub fn quote_style(&self) -> Quote {
720        FStringPartRef::from(self).quote_style()
721    }
722
723    pub fn is_empty_literal(&self) -> bool {
724        match &self {
725            FStringPart::Literal(string_literal) => string_literal.value.is_empty(),
726            FStringPart::FString(f_string) => f_string.elements.is_empty(),
727        }
728    }
729}
730
731impl Ranged for FStringPart {
732    fn range(&self) -> TextRange {
733        FStringPartRef::from(self).range()
734    }
735}
736
737/// A borrowed string literal or f-string part of an [`FStringValue`].
738#[derive(Copy, Clone, Debug, PartialEq)]
739pub enum FStringPartRef<'a> {
740    Literal(&'a StringLiteral),
741    FString(&'a FString),
742}
743
744impl FStringPartRef<'_> {
745    pub fn quote_style(self) -> Quote {
746        match self {
747            Self::Literal(literal) => literal.flags.quote_style(),
748            Self::FString(fstring) => fstring.flags.quote_style(),
749        }
750    }
751}
752
753impl<'a> From<&'a FStringPart> for FStringPartRef<'a> {
754    fn from(part: &'a FStringPart) -> Self {
755        match part {
756            FStringPart::Literal(literal) => Self::Literal(literal),
757            FStringPart::FString(fstring) => Self::FString(fstring),
758        }
759    }
760}
761
762impl Ranged for FStringPartRef<'_> {
763    fn range(&self) -> TextRange {
764        match self {
765            Self::Literal(literal) => literal.range(),
766            Self::FString(fstring) => fstring.range(),
767        }
768    }
769}
770
771/// A mutably borrowed string literal or f-string part of an [`FStringValue`].
772pub enum FStringPartMut<'a> {
773    Literal(&'a mut StringLiteral),
774    FString(&'a mut FString),
775}
776
777impl<'a> From<&'a mut FStringPart> for FStringPartMut<'a> {
778    fn from(part: &'a mut FStringPart) -> Self {
779        match part {
780            FStringPart::Literal(literal) => Self::Literal(literal),
781            FStringPart::FString(fstring) => Self::FString(fstring),
782        }
783    }
784}
785
786/// The iterator returned by [`FStringValue::iter`].
787#[derive(Clone)]
788pub enum FStringParts<'a> {
789    Single(Once<&'a FString>),
790    Concatenated(Iter<'a, FStringPart>),
791}
792
793impl<'a> Iterator for FStringParts<'a> {
794    type Item = FStringPartRef<'a>;
795
796    fn next(&mut self) -> Option<Self::Item> {
797        match self {
798            Self::Single(single) => single.next().map(FStringPartRef::FString),
799            Self::Concatenated(parts) => parts.next().map(FStringPartRef::from),
800        }
801    }
802
803    fn size_hint(&self) -> (usize, Option<usize>) {
804        match self {
805            Self::Single(single) => single.size_hint(),
806            Self::Concatenated(parts) => parts.size_hint(),
807        }
808    }
809}
810
811impl DoubleEndedIterator for FStringParts<'_> {
812    fn next_back(&mut self) -> Option<Self::Item> {
813        match self {
814            Self::Single(single) => single.next_back().map(FStringPartRef::FString),
815            Self::Concatenated(parts) => parts.next_back().map(FStringPartRef::from),
816        }
817    }
818}
819
820impl ExactSizeIterator for FStringParts<'_> {}
821impl FusedIterator for FStringParts<'_> {}
822
823/// The iterator returned by [`FStringValue::iter_mut`].
824pub enum FStringPartsMut<'a> {
825    Single(Once<&'a mut FString>),
826    Concatenated(IterMut<'a, FStringPart>),
827}
828
829impl<'a> Iterator for FStringPartsMut<'a> {
830    type Item = FStringPartMut<'a>;
831
832    fn next(&mut self) -> Option<Self::Item> {
833        match self {
834            Self::Single(single) => single.next().map(FStringPartMut::FString),
835            Self::Concatenated(parts) => parts.next().map(FStringPartMut::from),
836        }
837    }
838
839    fn size_hint(&self) -> (usize, Option<usize>) {
840        match self {
841            Self::Single(single) => single.size_hint(),
842            Self::Concatenated(parts) => parts.size_hint(),
843        }
844    }
845}
846
847impl DoubleEndedIterator for FStringPartsMut<'_> {
848    fn next_back(&mut self) -> Option<Self::Item> {
849        match self {
850            Self::Single(single) => single.next_back().map(FStringPartMut::FString),
851            Self::Concatenated(parts) => parts.next_back().map(FStringPartMut::from),
852        }
853    }
854}
855
856impl ExactSizeIterator for FStringPartsMut<'_> {}
857impl FusedIterator for FStringPartsMut<'_> {}
858
859impl ExprTString {
860    /// Returns the single [`TString`] if the t-string isn't implicitly concatenated, [`None`]
861    /// otherwise.
862    pub const fn as_single_part_tstring(&self) -> Option<&TString> {
863        match &self.value.inner {
864            TStringValueInner::Single(tstring) => Some(tstring),
865            TStringValueInner::Concatenated(_) => None,
866        }
867    }
868}
869
870/// The value representing an [`ExprTString`].
871#[derive(Clone, Debug, PartialEq)]
872#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
873pub struct TStringValue {
874    inner: TStringValueInner,
875}
876
877impl TStringValue {
878    /// Creates a new t-string literal with a single [`TString`] part.
879    pub fn single(value: TString) -> Self {
880        Self {
881            inner: TStringValueInner::Single(value),
882        }
883    }
884
885    /// Creates a new t-string with the given values that represents an implicitly
886    /// concatenated t-string.
887    ///
888    /// # Panics
889    ///
890    /// Panics if `values` has less than 2 elements.
891    /// Use [`TStringValue::single`] instead.
892    pub fn concatenated(values: Vec<TString>) -> Self {
893        assert!(
894            values.len() > 1,
895            "Use `TStringValue::single` to create single-part t-strings"
896        );
897        Self {
898            inner: TStringValueInner::Concatenated(values),
899        }
900    }
901
902    /// Returns `true` if the t-string is implicitly concatenated, `false` otherwise.
903    pub fn is_implicit_concatenated(&self) -> bool {
904        matches!(self.inner, TStringValueInner::Concatenated(_))
905    }
906
907    /// Returns a slice of all the [`TString`]s contained in this value.
908    pub fn as_slice(&self) -> &[TString] {
909        match &self.inner {
910            TStringValueInner::Single(part) => std::slice::from_ref(part),
911            TStringValueInner::Concatenated(parts) => parts,
912        }
913    }
914
915    /// Returns a mutable slice of all the [`TString`]s contained in this value.
916    fn as_mut_slice(&mut self) -> &mut [TString] {
917        match &mut self.inner {
918            TStringValueInner::Single(part) => std::slice::from_mut(part),
919            TStringValueInner::Concatenated(parts) => parts,
920        }
921    }
922
923    /// Returns an iterator over all the [`TString`]s contained in this value.
924    pub fn iter(&self) -> Iter<'_, TString> {
925        self.as_slice().iter()
926    }
927
928    /// Returns an iterator over all the [`TString`]s contained in this value
929    /// that allows modification.
930    pub fn iter_mut(&mut self) -> IterMut<'_, TString> {
931        self.as_mut_slice().iter_mut()
932    }
933
934    /// Returns an iterator over all the [`InterpolatedStringElement`] contained in this value.
935    ///
936    /// An interpolated string element is what makes up an [`TString`] i.e., it is either a
937    /// string literal or an interpolation. In the following example,
938    ///
939    /// ```python
940    /// t"foo" t"bar {x}" t"baz" t"qux"
941    /// ```
942    ///
943    /// The interpolated string elements returned would be string literal (`"bar "`),
944    /// interpolation (`x`) and string literal (`"qux"`).
945    pub fn elements(&self) -> impl Iterator<Item = &InterpolatedStringElement> {
946        self.iter().flat_map(|tstring| tstring.elements.iter())
947    }
948
949    /// Returns `true` if the node represents an empty t-string in the
950    /// sense that `__iter__` returns an empty iterable.
951    ///
952    /// Beware that empty t-strings are still truthy, i.e. `bool(t"") == True`.
953    ///
954    /// Note that a [`TStringValue`] node will always contain at least one
955    /// [`TString`] node. This method checks whether each of the constituent
956    /// t-strings (in an implicitly concatenated t-string) are empty
957    /// in the above sense.
958    pub fn is_empty_iterable(&self) -> bool {
959        match &self.inner {
960            TStringValueInner::Single(tstring) => tstring.is_empty(),
961            TStringValueInner::Concatenated(tstrings) => tstrings.iter().all(TString::is_empty),
962        }
963    }
964}
965
966impl<'a> IntoIterator for &'a TStringValue {
967    type Item = &'a TString;
968    type IntoIter = Iter<'a, TString>;
969
970    fn into_iter(self) -> Self::IntoIter {
971        self.iter()
972    }
973}
974
975impl<'a> IntoIterator for &'a mut TStringValue {
976    type Item = &'a mut TString;
977    type IntoIter = IterMut<'a, TString>;
978    fn into_iter(self) -> Self::IntoIter {
979        self.iter_mut()
980    }
981}
982
983/// An internal representation of [`TStringValue`].
984#[derive(Clone, Debug, PartialEq)]
985#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
986enum TStringValueInner {
987    /// A single t-string i.e., `t"foo"`.
988    Single(TString),
989
990    /// An implicitly concatenated t-string i.e., `t"foo" t"bar {x}"`.
991    Concatenated(Vec<TString>),
992}
993
994pub trait StringFlags: Copy {
995    /// Does the string use single or double quotes in its opener and closer?
996    fn quote_style(self) -> Quote;
997
998    fn triple_quotes(self) -> TripleQuotes;
999
1000    fn prefix(self) -> AnyStringPrefix;
1001
1002    fn is_unclosed(self) -> bool;
1003
1004    /// Is the string triple-quoted, i.e.,
1005    /// does it begin and end with three consecutive quote characters?
1006    fn is_triple_quoted(self) -> bool {
1007        self.triple_quotes().is_yes()
1008    }
1009
1010    /// A `str` representation of the quotes used to start and close.
1011    /// This does not include any prefixes the string has in its opener.
1012    fn quote_str(self) -> &'static str {
1013        match (self.triple_quotes(), self.quote_style()) {
1014            (TripleQuotes::Yes, Quote::Single) => "'''",
1015            (TripleQuotes::Yes, Quote::Double) => r#"""""#,
1016            (TripleQuotes::No, Quote::Single) => "'",
1017            (TripleQuotes::No, Quote::Double) => "\"",
1018        }
1019    }
1020
1021    /// The length of the quotes used to start and close the string.
1022    /// This does not include the length of any prefixes the string has
1023    /// in its opener.
1024    fn quote_len(self) -> TextSize {
1025        if self.is_triple_quoted() {
1026            TextSize::new(3)
1027        } else {
1028            TextSize::new(1)
1029        }
1030    }
1031
1032    /// The total length of the string's opener,
1033    /// i.e., the length of the prefixes plus the length
1034    /// of the quotes used to open the string.
1035    fn opener_len(self) -> TextSize {
1036        self.prefix().text_len() + self.quote_len()
1037    }
1038
1039    /// The total length of the string's closer.
1040    /// This is always equal to `self.quote_len()`, except when the string is unclosed,
1041    /// in which case the length is zero.
1042    fn closer_len(self) -> TextSize {
1043        if self.is_unclosed() {
1044            TextSize::default()
1045        } else {
1046            self.quote_len()
1047        }
1048    }
1049
1050    fn as_any_string_flags(self) -> AnyStringFlags {
1051        AnyStringFlags::new(self.prefix(), self.quote_style(), self.triple_quotes())
1052            .with_unclosed(self.is_unclosed())
1053    }
1054
1055    fn display_contents(self, contents: &str) -> DisplayFlags<'_> {
1056        DisplayFlags {
1057            flags: self.as_any_string_flags(),
1058            contents,
1059        }
1060    }
1061}
1062
1063pub struct DisplayFlags<'a> {
1064    flags: AnyStringFlags,
1065    contents: &'a str,
1066}
1067
1068impl std::fmt::Display for DisplayFlags<'_> {
1069    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1070        write!(
1071            f,
1072            "{prefix}{quote}{contents}{quote}",
1073            prefix = self.flags.prefix(),
1074            quote = self.flags.quote_str(),
1075            contents = self.contents
1076        )
1077    }
1078}
1079
1080bitflags! {
1081    #[derive(Default, Copy, Clone, PartialEq, Eq, Hash)]
1082    struct InterpolatedStringFlagsInner: u8 {
1083        /// The f-string uses double quotes (`"`) for its opener and closer.
1084        /// If this flag is not set, the f-string uses single quotes (`'`)
1085        /// for its opener and closer.
1086        const DOUBLE = 1 << 0;
1087
1088        /// The f-string is triple-quoted:
1089        /// it begins and ends with three consecutive quote characters.
1090        /// For example: `f"""{bar}"""`.
1091        const TRIPLE_QUOTED = 1 << 1;
1092
1093        /// The f-string has an `r` prefix, meaning it is a raw f-string
1094        /// with a lowercase 'r'. For example: `rf"{bar}"`
1095        const R_PREFIX_LOWER = 1 << 2;
1096
1097        /// The f-string has an `R` prefix, meaning it is a raw f-string
1098        /// with an uppercase 'r'. For example: `Rf"{bar}"`.
1099        /// See https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#r-strings-and-r-strings
1100        /// for why we track the casing of the `r` prefix,
1101        /// but not for any other prefix
1102        const R_PREFIX_UPPER = 1 << 3;
1103
1104        /// The f-string is unclosed, meaning it is missing a closing quote.
1105        /// For example: `f"{bar`
1106        const UNCLOSED = 1 << 4;
1107    }
1108}
1109
1110#[cfg(feature = "get-size")]
1111impl get_size2::GetSize for InterpolatedStringFlagsInner {}
1112
1113/// Flags that can be queried to obtain information
1114/// regarding the prefixes and quotes used for an f-string.
1115///
1116/// Note: This is identical to [`TStringFlags`] except that
1117/// the implementation of the `prefix` method of the
1118/// [`StringFlags`] trait returns a variant of
1119/// `AnyStringPrefix::Format`.
1120///
1121/// ## Notes on usage
1122///
1123/// If you're using a `Generator` from the `ruff_python_codegen` crate to generate a lint-rule fix
1124/// from an existing f-string literal, consider passing along the [`FString::flags`] field. If you
1125/// don't have an existing literal but have a `Checker` from the `ruff_linter` crate available,
1126/// consider using `Checker::default_fstring_flags` to create instances of this struct; this method
1127/// will properly handle nested f-strings. For usage that doesn't fit into one of these categories,
1128/// the public constructor [`FStringFlags::empty`] can be used.
1129#[derive(Copy, Clone, Eq, PartialEq, Hash)]
1130#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
1131pub struct FStringFlags(InterpolatedStringFlagsInner);
1132
1133impl FStringFlags {
1134    /// Construct a new [`FStringFlags`] with **no flags set**.
1135    ///
1136    /// See [`FStringFlags::with_quote_style`], [`FStringFlags::with_triple_quotes`], and
1137    /// [`FStringFlags::with_prefix`] for ways of setting the quote style (single or double),
1138    /// enabling triple quotes, and adding prefixes (such as `r`), respectively.
1139    ///
1140    /// See the documentation for [`FStringFlags`] for additional caveats on this constructor, and
1141    /// situations in which alternative ways to construct this struct should be used, especially
1142    /// when writing lint rules.
1143    pub fn empty() -> Self {
1144        Self(InterpolatedStringFlagsInner::empty())
1145    }
1146
1147    #[must_use]
1148    pub fn with_quote_style(mut self, quote_style: Quote) -> Self {
1149        self.0.set(
1150            InterpolatedStringFlagsInner::DOUBLE,
1151            quote_style.is_double(),
1152        );
1153        self
1154    }
1155
1156    #[must_use]
1157    pub fn with_triple_quotes(mut self, triple_quotes: TripleQuotes) -> Self {
1158        self.0.set(
1159            InterpolatedStringFlagsInner::TRIPLE_QUOTED,
1160            triple_quotes.is_yes(),
1161        );
1162        self
1163    }
1164
1165    #[must_use]
1166    pub fn with_unclosed(mut self, unclosed: bool) -> Self {
1167        self.0.set(InterpolatedStringFlagsInner::UNCLOSED, unclosed);
1168        self
1169    }
1170
1171    #[must_use]
1172    pub fn with_prefix(mut self, prefix: FStringPrefix) -> Self {
1173        match prefix {
1174            FStringPrefix::Regular => Self(
1175                self.0
1176                    - InterpolatedStringFlagsInner::R_PREFIX_LOWER
1177                    - InterpolatedStringFlagsInner::R_PREFIX_UPPER,
1178            ),
1179            FStringPrefix::Raw { uppercase_r } => {
1180                self.0
1181                    .set(InterpolatedStringFlagsInner::R_PREFIX_UPPER, uppercase_r);
1182                self.0
1183                    .set(InterpolatedStringFlagsInner::R_PREFIX_LOWER, !uppercase_r);
1184                self
1185            }
1186        }
1187    }
1188
1189    pub const fn prefix(self) -> FStringPrefix {
1190        if self
1191            .0
1192            .contains(InterpolatedStringFlagsInner::R_PREFIX_LOWER)
1193        {
1194            debug_assert!(
1195                !self
1196                    .0
1197                    .contains(InterpolatedStringFlagsInner::R_PREFIX_UPPER)
1198            );
1199            FStringPrefix::Raw { uppercase_r: false }
1200        } else if self
1201            .0
1202            .contains(InterpolatedStringFlagsInner::R_PREFIX_UPPER)
1203        {
1204            FStringPrefix::Raw { uppercase_r: true }
1205        } else {
1206            FStringPrefix::Regular
1207        }
1208    }
1209}
1210
1211// TODO(dylan): the documentation about using
1212// `Checker::default_tstring_flags` is not yet
1213// correct. This method does not yet exist because
1214// introducing it would emit a dead code warning
1215// until we call it in lint rules.
1216/// Flags that can be queried to obtain information
1217/// regarding the prefixes and quotes used for an f-string.
1218///
1219/// Note: This is identical to [`FStringFlags`] except that
1220/// the implementation of the `prefix` method of the
1221/// [`StringFlags`] trait returns a variant of
1222/// `AnyStringPrefix::Template`.
1223///
1224/// ## Notes on usage
1225///
1226/// If you're using a `Generator` from the `ruff_python_codegen` crate to generate a lint-rule fix
1227/// from an existing t-string literal, consider passing along the [`FString::flags`] field. If you
1228/// don't have an existing literal but have a `Checker` from the `ruff_linter` crate available,
1229/// consider using `Checker::default_tstring_flags` to create instances of this struct; this method
1230/// will properly handle nested t-strings. For usage that doesn't fit into one of these categories,
1231/// the public constructor [`TStringFlags::empty`] can be used.
1232#[derive(Copy, Clone, Eq, PartialEq, Hash)]
1233#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
1234pub struct TStringFlags(InterpolatedStringFlagsInner);
1235
1236impl TStringFlags {
1237    /// Construct a new [`TStringFlags`] with **no flags set**.
1238    ///
1239    /// See [`TStringFlags::with_quote_style`], [`TStringFlags::with_triple_quotes`], and
1240    /// [`TStringFlags::with_prefix`] for ways of setting the quote style (single or double),
1241    /// enabling triple quotes, and adding prefixes (such as `r`), respectively.
1242    ///
1243    /// See the documentation for [`TStringFlags`] for additional caveats on this constructor, and
1244    /// situations in which alternative ways to construct this struct should be used, especially
1245    /// when writing lint rules.
1246    pub fn empty() -> Self {
1247        Self(InterpolatedStringFlagsInner::empty())
1248    }
1249
1250    #[must_use]
1251    pub fn with_quote_style(mut self, quote_style: Quote) -> Self {
1252        self.0.set(
1253            InterpolatedStringFlagsInner::DOUBLE,
1254            quote_style.is_double(),
1255        );
1256        self
1257    }
1258
1259    #[must_use]
1260    pub fn with_triple_quotes(mut self, triple_quotes: TripleQuotes) -> Self {
1261        self.0.set(
1262            InterpolatedStringFlagsInner::TRIPLE_QUOTED,
1263            triple_quotes.is_yes(),
1264        );
1265        self
1266    }
1267
1268    #[must_use]
1269    pub fn with_unclosed(mut self, unclosed: bool) -> Self {
1270        self.0.set(InterpolatedStringFlagsInner::UNCLOSED, unclosed);
1271        self
1272    }
1273
1274    #[must_use]
1275    pub fn with_prefix(mut self, prefix: TStringPrefix) -> Self {
1276        match prefix {
1277            TStringPrefix::Regular => Self(
1278                self.0
1279                    - InterpolatedStringFlagsInner::R_PREFIX_LOWER
1280                    - InterpolatedStringFlagsInner::R_PREFIX_UPPER,
1281            ),
1282            TStringPrefix::Raw { uppercase_r } => {
1283                self.0
1284                    .set(InterpolatedStringFlagsInner::R_PREFIX_UPPER, uppercase_r);
1285                self.0
1286                    .set(InterpolatedStringFlagsInner::R_PREFIX_LOWER, !uppercase_r);
1287                self
1288            }
1289        }
1290    }
1291
1292    pub const fn prefix(self) -> TStringPrefix {
1293        if self
1294            .0
1295            .contains(InterpolatedStringFlagsInner::R_PREFIX_LOWER)
1296        {
1297            debug_assert!(
1298                !self
1299                    .0
1300                    .contains(InterpolatedStringFlagsInner::R_PREFIX_UPPER)
1301            );
1302            TStringPrefix::Raw { uppercase_r: false }
1303        } else if self
1304            .0
1305            .contains(InterpolatedStringFlagsInner::R_PREFIX_UPPER)
1306        {
1307            TStringPrefix::Raw { uppercase_r: true }
1308        } else {
1309            TStringPrefix::Regular
1310        }
1311    }
1312}
1313
1314impl StringFlags for FStringFlags {
1315    /// Return `true` if the f-string is triple-quoted, i.e.,
1316    /// it begins and ends with three consecutive quote characters.
1317    /// For example: `f"""{bar}"""`
1318    fn triple_quotes(self) -> TripleQuotes {
1319        if self.0.contains(InterpolatedStringFlagsInner::TRIPLE_QUOTED) {
1320            TripleQuotes::Yes
1321        } else {
1322            TripleQuotes::No
1323        }
1324    }
1325
1326    /// Return the quoting style (single or double quotes)
1327    /// used by the f-string's opener and closer:
1328    /// - `f"{"a"}"` -> `QuoteStyle::Double`
1329    /// - `f'{"a"}'` -> `QuoteStyle::Single`
1330    fn quote_style(self) -> Quote {
1331        if self.0.contains(InterpolatedStringFlagsInner::DOUBLE) {
1332            Quote::Double
1333        } else {
1334            Quote::Single
1335        }
1336    }
1337
1338    fn prefix(self) -> AnyStringPrefix {
1339        AnyStringPrefix::Format(self.prefix())
1340    }
1341
1342    fn is_unclosed(self) -> bool {
1343        self.0.intersects(InterpolatedStringFlagsInner::UNCLOSED)
1344    }
1345}
1346
1347impl fmt::Debug for FStringFlags {
1348    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1349        f.debug_struct("FStringFlags")
1350            .field("quote_style", &self.quote_style())
1351            .field("prefix", &self.prefix())
1352            .field("triple_quoted", &self.is_triple_quoted())
1353            .field("unclosed", &self.is_unclosed())
1354            .finish()
1355    }
1356}
1357
1358impl StringFlags for TStringFlags {
1359    /// Return `true` if the t-string is triple-quoted, i.e.,
1360    /// it begins and ends with three consecutive quote characters.
1361    /// For example: `t"""{bar}"""`
1362    fn triple_quotes(self) -> TripleQuotes {
1363        if self.0.contains(InterpolatedStringFlagsInner::TRIPLE_QUOTED) {
1364            TripleQuotes::Yes
1365        } else {
1366            TripleQuotes::No
1367        }
1368    }
1369
1370    /// Return the quoting style (single or double quotes)
1371    /// used by the t-string's opener and closer:
1372    /// - `t"{"a"}"` -> `QuoteStyle::Double`
1373    /// - `t'{"a"}'` -> `QuoteStyle::Single`
1374    fn quote_style(self) -> Quote {
1375        if self.0.contains(InterpolatedStringFlagsInner::DOUBLE) {
1376            Quote::Double
1377        } else {
1378            Quote::Single
1379        }
1380    }
1381
1382    fn prefix(self) -> AnyStringPrefix {
1383        AnyStringPrefix::Template(self.prefix())
1384    }
1385
1386    fn is_unclosed(self) -> bool {
1387        self.0.intersects(InterpolatedStringFlagsInner::UNCLOSED)
1388    }
1389}
1390
1391impl fmt::Debug for TStringFlags {
1392    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1393        f.debug_struct("TStringFlags")
1394            .field("quote_style", &self.quote_style())
1395            .field("prefix", &self.prefix())
1396            .field("triple_quoted", &self.is_triple_quoted())
1397            .field("unclosed", &self.is_unclosed())
1398            .finish()
1399    }
1400}
1401
1402/// An AST node that represents a single f-string which is part of an [`ExprFString`].
1403#[derive(Clone, Debug, PartialEq)]
1404#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
1405pub struct FString {
1406    pub range: TextRange,
1407    pub node_index: AtomicNodeIndex,
1408    pub elements: InterpolatedStringElements,
1409    pub flags: FStringFlags,
1410}
1411
1412impl From<FString> for Expr {
1413    fn from(payload: FString) -> Self {
1414        ExprFString {
1415            node_index: payload.node_index.clone(),
1416            range: payload.range,
1417            value: FStringValue::single(payload),
1418            runtime_joined_str: None,
1419            runtime_values: None,
1420        }
1421        .into()
1422    }
1423}
1424
1425/// A newtype wrapper around a list of [`InterpolatedStringElement`].
1426#[derive(Clone, Default, PartialEq)]
1427#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
1428pub struct InterpolatedStringElements(Box<[InterpolatedStringElement]>);
1429
1430impl InterpolatedStringElements {
1431    /// Returns an iterator over all the [`InterpolatedStringLiteralElement`] nodes contained in this f-string.
1432    pub fn literals(&self) -> impl Iterator<Item = &InterpolatedStringLiteralElement> {
1433        self.iter().filter_map(|element| element.as_literal())
1434    }
1435
1436    /// Returns an iterator over all the [`InterpolatedElement`] nodes contained in this f-string.
1437    pub fn interpolations(&self) -> impl Iterator<Item = &InterpolatedElement> {
1438        self.iter().filter_map(|element| element.as_interpolation())
1439    }
1440}
1441
1442impl From<Vec<InterpolatedStringElement>> for InterpolatedStringElements {
1443    fn from(elements: Vec<InterpolatedStringElement>) -> Self {
1444        InterpolatedStringElements(elements.into_boxed_slice())
1445    }
1446}
1447
1448impl<'a> IntoIterator for &'a InterpolatedStringElements {
1449    type IntoIter = Iter<'a, InterpolatedStringElement>;
1450    type Item = &'a InterpolatedStringElement;
1451
1452    fn into_iter(self) -> Self::IntoIter {
1453        self.iter()
1454    }
1455}
1456
1457impl<'a> IntoIterator for &'a mut InterpolatedStringElements {
1458    type IntoIter = IterMut<'a, InterpolatedStringElement>;
1459    type Item = &'a mut InterpolatedStringElement;
1460
1461    fn into_iter(self) -> Self::IntoIter {
1462        self.iter_mut()
1463    }
1464}
1465
1466impl Deref for InterpolatedStringElements {
1467    type Target = [InterpolatedStringElement];
1468
1469    fn deref(&self) -> &Self::Target {
1470        &self.0
1471    }
1472}
1473
1474impl DerefMut for InterpolatedStringElements {
1475    fn deref_mut(&mut self) -> &mut Self::Target {
1476        &mut self.0
1477    }
1478}
1479
1480impl fmt::Debug for InterpolatedStringElements {
1481    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1482        fmt::Debug::fmt(&self.0, f)
1483    }
1484}
1485
1486/// An AST node that represents a single t-string which is part of an [`ExprTString`].
1487#[derive(Clone, Debug, PartialEq)]
1488#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
1489pub struct TString {
1490    pub range: TextRange,
1491    pub node_index: AtomicNodeIndex,
1492    pub elements: InterpolatedStringElements,
1493    pub flags: TStringFlags,
1494}
1495
1496impl TString {
1497    pub fn quote_style(&self) -> Quote {
1498        self.flags.quote_style()
1499    }
1500
1501    pub fn is_empty(&self) -> bool {
1502        self.elements.is_empty()
1503    }
1504}
1505
1506impl From<TString> for Expr {
1507    fn from(payload: TString) -> Self {
1508        ExprTString {
1509            node_index: payload.node_index.clone(),
1510            range: payload.range,
1511            value: TStringValue::single(payload),
1512            runtime_template_str: None,
1513            runtime_values: None,
1514        }
1515        .into()
1516    }
1517}
1518
1519impl ExprStringLiteral {
1520    /// Return `Some(literal)` if the string only consists of a single `StringLiteral` part
1521    /// (indicating that it is not implicitly concatenated). Otherwise, return `None`.
1522    pub fn as_single_part_string(&self) -> Option<&StringLiteral> {
1523        match &self.value.inner {
1524            StringLiteralValueInner::Single(value) => Some(value),
1525            StringLiteralValueInner::Concatenated(_) => None,
1526        }
1527    }
1528}
1529
1530impl Ranged for ExprCall {
1531    fn range(&self) -> TextRange {
1532        TextRange::new(self.range_start, self.arguments.end())
1533    }
1534}
1535
1536#[expect(
1537    clippy::missing_fields_in_debug,
1538    reason = "`range_start` is represented by the reconstructed `range` field"
1539)]
1540impl fmt::Debug for ExprCall {
1541    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1542        formatter
1543            .debug_struct("ExprCall")
1544            .field("node_index", &self.node_index)
1545            .field("range", &self.range())
1546            .field("func", &self.func)
1547            .field("arguments", &self.arguments)
1548            .finish()
1549    }
1550}
1551
1552/// The value representing a [`ExprStringLiteral`].
1553#[derive(Clone, Debug, PartialEq)]
1554#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
1555pub struct StringLiteralValue {
1556    inner: StringLiteralValueInner,
1557}
1558
1559impl StringLiteralValue {
1560    /// Creates a new string literal with a single [`StringLiteral`] part.
1561    pub fn single(string: StringLiteral) -> Self {
1562        Self {
1563            inner: StringLiteralValueInner::Single(string),
1564        }
1565    }
1566
1567    /// Returns the [`StringLiteralFlags`] associated with this string literal.
1568    ///
1569    /// For an implicitly concatenated string, it returns the flags for the first literal.
1570    pub fn first_literal_flags(&self) -> StringLiteralFlags {
1571        self.iter()
1572            .next()
1573            .expect(
1574                "There should always be at least one string literal in an `ExprStringLiteral` node",
1575            )
1576            .flags
1577    }
1578
1579    /// Creates a new string literal with the given values that represents an
1580    /// implicitly concatenated strings.
1581    ///
1582    /// # Panics
1583    ///
1584    /// Panics if `strings` has less than 2 elements.
1585    /// Use [`StringLiteralValue::single`] instead.
1586    pub fn concatenated(strings: Vec<StringLiteral>) -> Self {
1587        assert!(
1588            strings.len() > 1,
1589            "Use `StringLiteralValue::single` to create single-part strings"
1590        );
1591        Self {
1592            inner: StringLiteralValueInner::Concatenated(Box::new(ConcatenatedStringLiteral {
1593                strings,
1594                value: OnceLock::new(),
1595            })),
1596        }
1597    }
1598
1599    /// Returns `true` if the string literal is implicitly concatenated.
1600    pub const fn is_implicit_concatenated(&self) -> bool {
1601        matches!(self.inner, StringLiteralValueInner::Concatenated(_))
1602    }
1603
1604    /// Returns `true` if the string literal has a `u` prefix, e.g. `u"foo"`.
1605    ///
1606    /// Although all strings in Python 3 are valid unicode (and the `u` prefix
1607    /// is only retained for backwards compatibility), these strings are known as
1608    /// "unicode strings".
1609    ///
1610    /// For an implicitly concatenated string, it returns `true` only if the first
1611    /// [`StringLiteral`] has the `u` prefix.
1612    pub fn is_unicode(&self) -> bool {
1613        self.iter()
1614            .next()
1615            .is_some_and(|part| part.flags.prefix().is_unicode())
1616    }
1617
1618    /// Returns a slice of all the [`StringLiteral`] parts contained in this value.
1619    pub fn as_slice(&self) -> &[StringLiteral] {
1620        match &self.inner {
1621            StringLiteralValueInner::Single(value) => std::slice::from_ref(value),
1622            StringLiteralValueInner::Concatenated(value) => value.strings.as_slice(),
1623        }
1624    }
1625
1626    /// Returns a mutable slice of all the [`StringLiteral`] parts contained in this value.
1627    fn as_mut_slice(&mut self) -> &mut [StringLiteral] {
1628        match &mut self.inner {
1629            StringLiteralValueInner::Single(value) => std::slice::from_mut(value),
1630            StringLiteralValueInner::Concatenated(value) => value.strings.as_mut_slice(),
1631        }
1632    }
1633
1634    /// Returns an iterator over all the [`StringLiteral`] parts contained in this value.
1635    pub fn iter(&self) -> Iter<'_, StringLiteral> {
1636        self.as_slice().iter()
1637    }
1638
1639    /// Returns an iterator over all the [`StringLiteral`] parts contained in this value
1640    /// that allows modification.
1641    pub fn iter_mut(&mut self) -> IterMut<'_, StringLiteral> {
1642        self.as_mut_slice().iter_mut()
1643    }
1644
1645    /// Returns `true` if the node represents an empty string.
1646    ///
1647    /// Note that a [`StringLiteralValue`] node will always have >=1 [`StringLiteral`] parts
1648    /// inside it. This method checks whether the value of the concatenated parts is equal
1649    /// to the empty string, not whether the string has 0 parts inside it.
1650    pub fn is_empty(&self) -> bool {
1651        self.len() == 0
1652    }
1653
1654    /// Returns the total length of the string literal value, in bytes, not
1655    /// [`char`]s or graphemes.
1656    pub fn len(&self) -> usize {
1657        self.iter().fold(0, |acc, part| acc + part.value.len())
1658    }
1659
1660    /// Returns an iterator over the [`char`]s of each string literal part.
1661    pub fn chars(&self) -> impl Iterator<Item = char> + Clone + '_ {
1662        self.iter().flat_map(|part| part.value.chars())
1663    }
1664
1665    /// Returns the concatenated string value as a [`str`].
1666    ///
1667    /// Note that this will perform an allocation on the first invocation if the
1668    /// string value is implicitly concatenated.
1669    pub fn to_str(&self) -> &str {
1670        match &self.inner {
1671            StringLiteralValueInner::Single(value) => value.as_str(),
1672            StringLiteralValueInner::Concatenated(value) => value.to_str(),
1673        }
1674    }
1675}
1676
1677impl<'a> IntoIterator for &'a StringLiteralValue {
1678    type Item = &'a StringLiteral;
1679    type IntoIter = Iter<'a, StringLiteral>;
1680
1681    fn into_iter(self) -> Self::IntoIter {
1682        self.iter()
1683    }
1684}
1685
1686impl<'a> IntoIterator for &'a mut StringLiteralValue {
1687    type Item = &'a mut StringLiteral;
1688    type IntoIter = IterMut<'a, StringLiteral>;
1689    fn into_iter(self) -> Self::IntoIter {
1690        self.iter_mut()
1691    }
1692}
1693
1694impl PartialEq<str> for StringLiteralValue {
1695    fn eq(&self, other: &str) -> bool {
1696        if self.len() != other.len() {
1697            return false;
1698        }
1699        // The `zip` here is safe because we have checked the length of both parts.
1700        self.chars().zip(other.chars()).all(|(c1, c2)| c1 == c2)
1701    }
1702}
1703
1704impl fmt::Display for StringLiteralValue {
1705    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1706        f.write_str(self.to_str())
1707    }
1708}
1709
1710/// An internal representation of [`StringLiteralValue`].
1711#[derive(Clone, Debug, PartialEq)]
1712#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
1713enum StringLiteralValueInner {
1714    /// A single string literal i.e., `"foo"`.
1715    Single(StringLiteral),
1716
1717    /// An implicitly concatenated string literals i.e., `"foo" "bar"`.
1718    Concatenated(Box<ConcatenatedStringLiteral>),
1719}
1720
1721bitflags! {
1722    #[derive(Debug, Default, Copy, Clone, PartialEq, Eq, Hash)]
1723    struct StringLiteralFlagsInner: u8 {
1724        /// The string uses double quotes (e.g. `"foo"`).
1725        /// If this flag is not set, the string uses single quotes (`'foo'`).
1726        const DOUBLE = 1 << 0;
1727
1728        /// The string is triple-quoted (`"""foo"""`):
1729        /// it begins and ends with three consecutive quote characters.
1730        const TRIPLE_QUOTED = 1 << 1;
1731
1732        /// The string has a `u` or `U` prefix, e.g. `u"foo"`.
1733        /// While this prefix is a no-op at runtime,
1734        /// strings with this prefix can have no other prefixes set;
1735        /// it is therefore invalid for this flag to be set
1736        /// if `R_PREFIX` is also set.
1737        const U_PREFIX = 1 << 2;
1738
1739        /// The string has an `r` prefix, meaning it is a raw string
1740        /// with a lowercase 'r' (e.g. `r"foo\."`).
1741        /// It is invalid to set this flag if `U_PREFIX` is also set.
1742        const R_PREFIX_LOWER = 1 << 3;
1743
1744        /// The string has an `R` prefix, meaning it is a raw string
1745        /// with an uppercase 'R' (e.g. `R'foo\d'`).
1746        /// See https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#r-strings-and-r-strings
1747        /// for why we track the casing of the `r` prefix,
1748        /// but not for any other prefix
1749        const R_PREFIX_UPPER = 1 << 4;
1750
1751        /// The string was deemed invalid by the parser.
1752        const INVALID = 1 << 5;
1753
1754        /// The string literal misses the matching closing quote(s).
1755        const UNCLOSED = 1 << 6;
1756    }
1757}
1758
1759#[cfg(feature = "get-size")]
1760impl get_size2::GetSize for StringLiteralFlagsInner {}
1761
1762/// Flags that can be queried to obtain information
1763/// regarding the prefixes and quotes used for a string literal.
1764///
1765/// ## Notes on usage
1766///
1767/// If you're using a `Generator` from the `ruff_python_codegen` crate to generate a lint-rule fix
1768/// from an existing string literal, consider passing along the [`StringLiteral::flags`] field or
1769/// the result of the [`StringLiteralValue::first_literal_flags`] method. If you don't have an
1770/// existing string but have a `Checker` from the `ruff_linter` crate available, consider using
1771/// `Checker::default_string_flags` to create instances of this struct; this method will properly
1772/// handle surrounding f-strings. For usage that doesn't fit into one of these categories, the
1773/// public constructor [`StringLiteralFlags::empty`] can be used.
1774#[derive(Copy, Clone, Eq, PartialEq, Hash)]
1775#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
1776pub struct StringLiteralFlags(StringLiteralFlagsInner);
1777
1778impl StringLiteralFlags {
1779    /// Construct a new [`StringLiteralFlags`] with **no flags set**.
1780    ///
1781    /// See [`StringLiteralFlags::with_quote_style`], [`StringLiteralFlags::with_triple_quotes`],
1782    /// and [`StringLiteralFlags::with_prefix`] for ways of setting the quote style (single or
1783    /// double), enabling triple quotes, and adding prefixes (such as `r` or `u`), respectively.
1784    ///
1785    /// See the documentation for [`StringLiteralFlags`] for additional caveats on this constructor,
1786    /// and situations in which alternative ways to construct this struct should be used, especially
1787    /// when writing lint rules.
1788    pub fn empty() -> Self {
1789        Self(StringLiteralFlagsInner::empty())
1790    }
1791
1792    #[must_use]
1793    pub fn with_quote_style(mut self, quote_style: Quote) -> Self {
1794        self.0
1795            .set(StringLiteralFlagsInner::DOUBLE, quote_style.is_double());
1796        self
1797    }
1798
1799    #[must_use]
1800    pub fn with_triple_quotes(mut self, triple_quotes: TripleQuotes) -> Self {
1801        self.0.set(
1802            StringLiteralFlagsInner::TRIPLE_QUOTED,
1803            triple_quotes.is_yes(),
1804        );
1805        self
1806    }
1807
1808    #[must_use]
1809    pub fn with_unclosed(mut self, unclosed: bool) -> Self {
1810        self.0.set(StringLiteralFlagsInner::UNCLOSED, unclosed);
1811        self
1812    }
1813
1814    #[must_use]
1815    pub fn with_prefix(self, prefix: StringLiteralPrefix) -> Self {
1816        let StringLiteralFlags(flags) = self;
1817        match prefix {
1818            StringLiteralPrefix::Empty => Self(
1819                flags
1820                    - StringLiteralFlagsInner::R_PREFIX_LOWER
1821                    - StringLiteralFlagsInner::R_PREFIX_UPPER
1822                    - StringLiteralFlagsInner::U_PREFIX,
1823            ),
1824            StringLiteralPrefix::Raw { uppercase: false } => Self(
1825                (flags | StringLiteralFlagsInner::R_PREFIX_LOWER)
1826                    - StringLiteralFlagsInner::R_PREFIX_UPPER
1827                    - StringLiteralFlagsInner::U_PREFIX,
1828            ),
1829            StringLiteralPrefix::Raw { uppercase: true } => Self(
1830                (flags | StringLiteralFlagsInner::R_PREFIX_UPPER)
1831                    - StringLiteralFlagsInner::R_PREFIX_LOWER
1832                    - StringLiteralFlagsInner::U_PREFIX,
1833            ),
1834            StringLiteralPrefix::Unicode => Self(
1835                (flags | StringLiteralFlagsInner::U_PREFIX)
1836                    - StringLiteralFlagsInner::R_PREFIX_LOWER
1837                    - StringLiteralFlagsInner::R_PREFIX_UPPER,
1838            ),
1839        }
1840    }
1841
1842    #[must_use]
1843    pub fn with_invalid(mut self) -> Self {
1844        self.0 |= StringLiteralFlagsInner::INVALID;
1845        self
1846    }
1847
1848    /// Returns `true` if the parser deemed the string literal invalid.
1849    pub const fn is_invalid(self) -> bool {
1850        self.0.contains(StringLiteralFlagsInner::INVALID)
1851    }
1852
1853    pub const fn prefix(self) -> StringLiteralPrefix {
1854        if self.0.contains(StringLiteralFlagsInner::U_PREFIX) {
1855            debug_assert!(
1856                !self.0.intersects(
1857                    StringLiteralFlagsInner::R_PREFIX_LOWER
1858                        .union(StringLiteralFlagsInner::R_PREFIX_UPPER)
1859                )
1860            );
1861            StringLiteralPrefix::Unicode
1862        } else if self.0.contains(StringLiteralFlagsInner::R_PREFIX_LOWER) {
1863            debug_assert!(!self.0.contains(StringLiteralFlagsInner::R_PREFIX_UPPER));
1864            StringLiteralPrefix::Raw { uppercase: false }
1865        } else if self.0.contains(StringLiteralFlagsInner::R_PREFIX_UPPER) {
1866            StringLiteralPrefix::Raw { uppercase: true }
1867        } else {
1868            StringLiteralPrefix::Empty
1869        }
1870    }
1871}
1872
1873impl StringFlags for StringLiteralFlags {
1874    /// Return the quoting style (single or double quotes)
1875    /// used by the string's opener and closer:
1876    /// - `"a"` -> `QuoteStyle::Double`
1877    /// - `'a'` -> `QuoteStyle::Single`
1878    fn quote_style(self) -> Quote {
1879        if self.0.contains(StringLiteralFlagsInner::DOUBLE) {
1880            Quote::Double
1881        } else {
1882            Quote::Single
1883        }
1884    }
1885
1886    /// Return `true` if the string is triple-quoted, i.e.,
1887    /// it begins and ends with three consecutive quote characters.
1888    /// For example: `"""bar"""`
1889    fn triple_quotes(self) -> TripleQuotes {
1890        if self.0.contains(StringLiteralFlagsInner::TRIPLE_QUOTED) {
1891            TripleQuotes::Yes
1892        } else {
1893            TripleQuotes::No
1894        }
1895    }
1896
1897    fn prefix(self) -> AnyStringPrefix {
1898        AnyStringPrefix::Regular(self.prefix())
1899    }
1900
1901    fn is_unclosed(self) -> bool {
1902        self.0.intersects(StringLiteralFlagsInner::UNCLOSED)
1903    }
1904}
1905
1906impl fmt::Debug for StringLiteralFlags {
1907    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1908        f.debug_struct("StringLiteralFlags")
1909            .field("quote_style", &self.quote_style())
1910            .field("prefix", &self.prefix())
1911            .field("triple_quoted", &self.is_triple_quoted())
1912            .field("unclosed", &self.is_unclosed())
1913            .finish()
1914    }
1915}
1916
1917/// An AST node that represents a single string literal which is part of an
1918/// [`ExprStringLiteral`].
1919#[derive(Clone, Debug, PartialEq)]
1920#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
1921pub struct StringLiteral {
1922    pub range: TextRange,
1923    pub node_index: AtomicNodeIndex,
1924    pub value: Box<str>,
1925    pub flags: StringLiteralFlags,
1926}
1927
1928impl Deref for StringLiteral {
1929    type Target = str;
1930
1931    fn deref(&self) -> &Self::Target {
1932        &self.value
1933    }
1934}
1935
1936impl StringLiteral {
1937    /// Extracts a string slice containing the entire `String`.
1938    pub fn as_str(&self) -> &str {
1939        self
1940    }
1941
1942    /// Creates an invalid string literal with the given range.
1943    pub fn invalid(range: TextRange) -> Self {
1944        Self {
1945            range,
1946            value: "".into(),
1947            node_index: AtomicNodeIndex::NONE,
1948            flags: StringLiteralFlags::empty().with_invalid(),
1949        }
1950    }
1951
1952    /// The range of the string literal's contents.
1953    ///
1954    /// This excludes any prefixes, opening quotes or closing quotes.
1955    pub fn content_range(&self) -> TextRange {
1956        TextRange::new(
1957            self.start() + self.flags.opener_len(),
1958            self.end() - self.flags.closer_len(),
1959        )
1960    }
1961}
1962
1963impl From<StringLiteral> for Expr {
1964    fn from(payload: StringLiteral) -> Self {
1965        ExprStringLiteral {
1966            range: payload.range,
1967            node_index: AtomicNodeIndex::NONE,
1968            value: StringLiteralValue::single(payload),
1969        }
1970        .into()
1971    }
1972}
1973
1974/// An internal representation of [`StringLiteral`] that represents an
1975/// implicitly concatenated string.
1976#[derive(Clone)]
1977#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
1978struct ConcatenatedStringLiteral {
1979    /// The individual [`StringLiteral`] parts that make up the concatenated string.
1980    strings: Vec<StringLiteral>,
1981    /// The concatenated string value.
1982    value: OnceLock<Box<str>>,
1983}
1984
1985impl ConcatenatedStringLiteral {
1986    /// Extracts a string slice containing the entire concatenated string.
1987    fn to_str(&self) -> &str {
1988        self.value.get_or_init(|| {
1989            let concatenated: String = self.strings.iter().map(StringLiteral::as_str).collect();
1990            concatenated.into_boxed_str()
1991        })
1992    }
1993}
1994
1995impl PartialEq for ConcatenatedStringLiteral {
1996    fn eq(&self, other: &Self) -> bool {
1997        if self.strings.len() != other.strings.len() {
1998            return false;
1999        }
2000        // The `zip` here is safe because we have checked the length of both parts.
2001        self.strings
2002            .iter()
2003            .zip(&other.strings)
2004            .all(|(s1, s2)| s1 == s2)
2005    }
2006}
2007
2008impl Debug for ConcatenatedStringLiteral {
2009    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2010        f.debug_struct("ConcatenatedStringLiteral")
2011            .field("strings", &self.strings)
2012            .field("value", &self.to_str())
2013            .finish()
2014    }
2015}
2016
2017impl ExprBytesLiteral {
2018    /// Return `Some(literal)` if the bytestring only consists of a single `BytesLiteral` part
2019    /// (indicating that it is not implicitly concatenated). Otherwise, return `None`.
2020    pub const fn as_single_part_bytestring(&self) -> Option<&BytesLiteral> {
2021        match &self.value.inner {
2022            BytesLiteralValueInner::Single(value) => Some(value),
2023            BytesLiteralValueInner::Concatenated(_) => None,
2024        }
2025    }
2026}
2027
2028/// The value representing a [`ExprBytesLiteral`].
2029#[derive(Clone, Debug, PartialEq)]
2030#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
2031pub struct BytesLiteralValue {
2032    inner: BytesLiteralValueInner,
2033}
2034
2035impl BytesLiteralValue {
2036    /// Create a new bytestring literal with a single [`BytesLiteral`] part.
2037    pub fn single(value: BytesLiteral) -> Self {
2038        Self {
2039            inner: BytesLiteralValueInner::Single(value),
2040        }
2041    }
2042
2043    /// Creates a new bytestring literal with the given values that represents an
2044    /// implicitly concatenated bytestring.
2045    ///
2046    /// # Panics
2047    ///
2048    /// Panics if `values` has less than 2 elements.
2049    /// Use [`BytesLiteralValue::single`] instead.
2050    pub fn concatenated(values: Vec<BytesLiteral>) -> Self {
2051        assert!(
2052            values.len() > 1,
2053            "Use `BytesLiteralValue::single` to create single-part bytestrings"
2054        );
2055        Self {
2056            inner: BytesLiteralValueInner::Concatenated(values),
2057        }
2058    }
2059
2060    /// Returns `true` if the bytestring is implicitly concatenated.
2061    pub const fn is_implicit_concatenated(&self) -> bool {
2062        matches!(self.inner, BytesLiteralValueInner::Concatenated(_))
2063    }
2064
2065    /// Returns a slice of all the [`BytesLiteral`] parts contained in this value.
2066    pub fn as_slice(&self) -> &[BytesLiteral] {
2067        match &self.inner {
2068            BytesLiteralValueInner::Single(value) => std::slice::from_ref(value),
2069            BytesLiteralValueInner::Concatenated(value) => value.as_slice(),
2070        }
2071    }
2072
2073    /// Returns a mutable slice of all the [`BytesLiteral`] parts contained in this value.
2074    fn as_mut_slice(&mut self) -> &mut [BytesLiteral] {
2075        match &mut self.inner {
2076            BytesLiteralValueInner::Single(value) => std::slice::from_mut(value),
2077            BytesLiteralValueInner::Concatenated(value) => value.as_mut_slice(),
2078        }
2079    }
2080
2081    /// Returns an iterator over all the [`BytesLiteral`] parts contained in this value.
2082    pub fn iter(&self) -> Iter<'_, BytesLiteral> {
2083        self.as_slice().iter()
2084    }
2085
2086    /// Returns an iterator over all the [`BytesLiteral`] parts contained in this value
2087    /// that allows modification.
2088    pub fn iter_mut(&mut self) -> IterMut<'_, BytesLiteral> {
2089        self.as_mut_slice().iter_mut()
2090    }
2091
2092    /// Return `true` if the node represents an empty bytestring.
2093    ///
2094    /// Note that a [`BytesLiteralValue`] node will always have >=1 [`BytesLiteral`] parts
2095    /// inside it. This method checks whether the value of the concatenated parts is equal
2096    /// to the empty bytestring, not whether the bytestring has 0 parts inside it.
2097    pub fn is_empty(&self) -> bool {
2098        self.iter().all(|part| part.is_empty())
2099    }
2100
2101    /// Returns the length of the concatenated bytestring.
2102    pub fn len(&self) -> usize {
2103        self.iter().map(|part| part.len()).sum()
2104    }
2105
2106    /// Returns an iterator over the bytes of the concatenated bytestring.
2107    pub fn bytes(&self) -> impl Iterator<Item = u8> + '_ {
2108        self.iter().flat_map(|part| part.as_slice().iter().copied())
2109    }
2110}
2111
2112impl<'a> IntoIterator for &'a BytesLiteralValue {
2113    type Item = &'a BytesLiteral;
2114    type IntoIter = Iter<'a, BytesLiteral>;
2115
2116    fn into_iter(self) -> Self::IntoIter {
2117        self.iter()
2118    }
2119}
2120
2121impl<'a> IntoIterator for &'a mut BytesLiteralValue {
2122    type Item = &'a mut BytesLiteral;
2123    type IntoIter = IterMut<'a, BytesLiteral>;
2124    fn into_iter(self) -> Self::IntoIter {
2125        self.iter_mut()
2126    }
2127}
2128
2129impl PartialEq<[u8]> for BytesLiteralValue {
2130    fn eq(&self, other: &[u8]) -> bool {
2131        if self.len() != other.len() {
2132            return false;
2133        }
2134        // The `zip` here is safe because we have checked the length of both parts.
2135        self.bytes()
2136            .zip(other.iter().copied())
2137            .all(|(b1, b2)| b1 == b2)
2138    }
2139}
2140
2141impl<'a> From<&'a BytesLiteralValue> for Cow<'a, [u8]> {
2142    fn from(value: &'a BytesLiteralValue) -> Self {
2143        match &value.inner {
2144            BytesLiteralValueInner::Single(BytesLiteral {
2145                value: bytes_value, ..
2146            }) => Cow::from(bytes_value.as_ref()),
2147            BytesLiteralValueInner::Concatenated(bytes_literal_vec) => Cow::Owned(
2148                bytes_literal_vec
2149                    .iter()
2150                    .flat_map(|bytes_literal| bytes_literal.value.to_vec())
2151                    .collect::<Vec<u8>>(),
2152            ),
2153        }
2154    }
2155}
2156
2157/// An internal representation of [`BytesLiteralValue`].
2158#[derive(Clone, Debug, PartialEq)]
2159#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
2160enum BytesLiteralValueInner {
2161    /// A single-part bytestring literal i.e., `b"foo"`.
2162    Single(BytesLiteral),
2163
2164    /// An implicitly concatenated bytestring literal i.e., `b"foo" b"bar"`.
2165    Concatenated(Vec<BytesLiteral>),
2166}
2167
2168bitflags! {
2169    #[derive(Default, Copy, Clone, PartialEq, Eq, Hash)]
2170    struct BytesLiteralFlagsInner: u8 {
2171        /// The bytestring uses double quotes (e.g. `b"foo"`).
2172        /// If this flag is not set, the bytestring uses single quotes (e.g. `b'foo'`).
2173        const DOUBLE = 1 << 0;
2174
2175        /// The bytestring is triple-quoted (e.g. `b"""foo"""`):
2176        /// it begins and ends with three consecutive quote characters.
2177        const TRIPLE_QUOTED = 1 << 1;
2178
2179        /// The bytestring has an `r` prefix (e.g. `rb"foo"`),
2180        /// meaning it is a raw bytestring with a lowercase 'r'.
2181        const R_PREFIX_LOWER = 1 << 2;
2182
2183        /// The bytestring has an `R` prefix (e.g. `Rb"foo"`),
2184        /// meaning it is a raw bytestring with an uppercase 'R'.
2185        /// See https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#r-strings-and-r-strings
2186        /// for why we track the casing of the `r` prefix, but not for any other prefix
2187        const R_PREFIX_UPPER = 1 << 3;
2188
2189        /// The bytestring was deemed invalid by the parser.
2190        const INVALID = 1 << 4;
2191
2192        /// The byte string misses the matching closing quote(s).
2193        const UNCLOSED = 1 << 5;
2194    }
2195}
2196
2197#[cfg(feature = "get-size")]
2198impl get_size2::GetSize for BytesLiteralFlagsInner {}
2199
2200/// Flags that can be queried to obtain information
2201/// regarding the prefixes and quotes used for a bytes literal.
2202///
2203/// ## Notes on usage
2204///
2205/// If you're using a `Generator` from the `ruff_python_codegen` crate to generate a lint-rule fix
2206/// from an existing bytes literal, consider passing along the [`BytesLiteral::flags`] field. If
2207/// you don't have an existing literal but have a `Checker` from the `ruff_linter` crate available,
2208/// consider using `Checker::default_bytes_flags` to create instances of this struct; this method
2209/// will properly handle surrounding f-strings. For usage that doesn't fit into one of these
2210/// categories, the public constructor [`BytesLiteralFlags::empty`] can be used.
2211#[derive(Copy, Clone, Eq, PartialEq, Hash)]
2212#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
2213pub struct BytesLiteralFlags(BytesLiteralFlagsInner);
2214
2215impl BytesLiteralFlags {
2216    /// Construct a new [`BytesLiteralFlags`] with **no flags set**.
2217    ///
2218    /// See [`BytesLiteralFlags::with_quote_style`], [`BytesLiteralFlags::with_triple_quotes`], and
2219    /// [`BytesLiteralFlags::with_prefix`] for ways of setting the quote style (single or double),
2220    /// enabling triple quotes, and adding prefixes (such as `r`), respectively.
2221    ///
2222    /// See the documentation for [`BytesLiteralFlags`] for additional caveats on this constructor,
2223    /// and situations in which alternative ways to construct this struct should be used, especially
2224    /// when writing lint rules.
2225    pub fn empty() -> Self {
2226        Self(BytesLiteralFlagsInner::empty())
2227    }
2228
2229    #[must_use]
2230    pub fn with_quote_style(mut self, quote_style: Quote) -> Self {
2231        self.0
2232            .set(BytesLiteralFlagsInner::DOUBLE, quote_style.is_double());
2233        self
2234    }
2235
2236    #[must_use]
2237    pub fn with_triple_quotes(mut self, triple_quotes: TripleQuotes) -> Self {
2238        self.0.set(
2239            BytesLiteralFlagsInner::TRIPLE_QUOTED,
2240            triple_quotes.is_yes(),
2241        );
2242        self
2243    }
2244
2245    #[must_use]
2246    pub fn with_unclosed(mut self, unclosed: bool) -> Self {
2247        self.0.set(BytesLiteralFlagsInner::UNCLOSED, unclosed);
2248        self
2249    }
2250
2251    #[must_use]
2252    pub fn with_prefix(mut self, prefix: ByteStringPrefix) -> Self {
2253        match prefix {
2254            ByteStringPrefix::Regular => {
2255                self.0 -= BytesLiteralFlagsInner::R_PREFIX_LOWER;
2256                self.0 -= BytesLiteralFlagsInner::R_PREFIX_UPPER;
2257            }
2258            ByteStringPrefix::Raw { uppercase_r } => {
2259                self.0
2260                    .set(BytesLiteralFlagsInner::R_PREFIX_UPPER, uppercase_r);
2261                self.0
2262                    .set(BytesLiteralFlagsInner::R_PREFIX_LOWER, !uppercase_r);
2263            }
2264        }
2265        self
2266    }
2267
2268    #[must_use]
2269    pub fn with_invalid(mut self) -> Self {
2270        self.0 |= BytesLiteralFlagsInner::INVALID;
2271        self
2272    }
2273
2274    /// Returns `true` if the parser deemed the bytes literal invalid.
2275    pub const fn is_invalid(self) -> bool {
2276        self.0.contains(BytesLiteralFlagsInner::INVALID)
2277    }
2278
2279    pub const fn prefix(self) -> ByteStringPrefix {
2280        if self.0.contains(BytesLiteralFlagsInner::R_PREFIX_LOWER) {
2281            debug_assert!(!self.0.contains(BytesLiteralFlagsInner::R_PREFIX_UPPER));
2282            ByteStringPrefix::Raw { uppercase_r: false }
2283        } else if self.0.contains(BytesLiteralFlagsInner::R_PREFIX_UPPER) {
2284            ByteStringPrefix::Raw { uppercase_r: true }
2285        } else {
2286            ByteStringPrefix::Regular
2287        }
2288    }
2289}
2290
2291impl StringFlags for BytesLiteralFlags {
2292    /// Return `true` if the bytestring is triple-quoted, i.e.,
2293    /// it begins and ends with three consecutive quote characters.
2294    /// For example: `b"""{bar}"""`
2295    fn triple_quotes(self) -> TripleQuotes {
2296        if self.0.contains(BytesLiteralFlagsInner::TRIPLE_QUOTED) {
2297            TripleQuotes::Yes
2298        } else {
2299            TripleQuotes::No
2300        }
2301    }
2302
2303    /// Return the quoting style (single or double quotes)
2304    /// used by the bytestring's opener and closer:
2305    /// - `b"a"` -> `QuoteStyle::Double`
2306    /// - `b'a'` -> `QuoteStyle::Single`
2307    fn quote_style(self) -> Quote {
2308        if self.0.contains(BytesLiteralFlagsInner::DOUBLE) {
2309            Quote::Double
2310        } else {
2311            Quote::Single
2312        }
2313    }
2314
2315    fn prefix(self) -> AnyStringPrefix {
2316        AnyStringPrefix::Bytes(self.prefix())
2317    }
2318
2319    fn is_unclosed(self) -> bool {
2320        self.0.intersects(BytesLiteralFlagsInner::UNCLOSED)
2321    }
2322}
2323
2324impl fmt::Debug for BytesLiteralFlags {
2325    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2326        f.debug_struct("BytesLiteralFlags")
2327            .field("quote_style", &self.quote_style())
2328            .field("prefix", &self.prefix())
2329            .field("triple_quoted", &self.is_triple_quoted())
2330            .field("unclosed", &self.is_unclosed())
2331            .finish()
2332    }
2333}
2334
2335/// An AST node that represents a single bytes literal which is part of an
2336/// [`ExprBytesLiteral`].
2337#[derive(Clone, Debug, PartialEq)]
2338#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
2339pub struct BytesLiteral {
2340    pub range: TextRange,
2341    pub node_index: AtomicNodeIndex,
2342    pub value: Box<[u8]>,
2343    pub flags: BytesLiteralFlags,
2344}
2345
2346impl Deref for BytesLiteral {
2347    type Target = [u8];
2348
2349    fn deref(&self) -> &Self::Target {
2350        &self.value
2351    }
2352}
2353
2354impl BytesLiteral {
2355    /// Extracts a byte slice containing the entire [`BytesLiteral`].
2356    pub fn as_slice(&self) -> &[u8] {
2357        self
2358    }
2359
2360    /// Creates a new invalid bytes literal with the given range.
2361    pub fn invalid(range: TextRange) -> Self {
2362        Self {
2363            range,
2364            value: Box::new([]),
2365            node_index: AtomicNodeIndex::NONE,
2366            flags: BytesLiteralFlags::empty().with_invalid(),
2367        }
2368    }
2369
2370    /// The range of the byte literal's contents.
2371    ///
2372    /// This excludes any prefixes, opening quotes or closing quotes.
2373    pub fn content_range(&self) -> TextRange {
2374        TextRange::new(
2375            self.start() + self.flags.opener_len(),
2376            self.end() - self.flags.closer_len(),
2377        )
2378    }
2379}
2380
2381impl From<BytesLiteral> for Expr {
2382    fn from(payload: BytesLiteral) -> Self {
2383        ExprBytesLiteral {
2384            range: payload.range,
2385            node_index: AtomicNodeIndex::NONE,
2386            value: BytesLiteralValue::single(payload),
2387        }
2388        .into()
2389    }
2390}
2391
2392bitflags! {
2393    /// Flags that can be queried to obtain information
2394    /// regarding the prefixes and quotes used for a string literal.
2395    ///
2396    /// Note that not all of these flags can be validly combined -- e.g.,
2397    /// it is invalid to combine the `U_PREFIX` flag with any other
2398    /// of the `*_PREFIX` flags. As such, the recommended way to set the
2399    /// prefix flags is by calling the `as_flags()` method on the
2400    /// `StringPrefix` enum.
2401    #[derive(Default, Debug, Copy, Clone, PartialEq, Eq, Hash)]
2402    struct AnyStringFlagsInner: u16 {
2403        /// The string uses double quotes (`"`).
2404        /// If this flag is not set, the string uses single quotes (`'`).
2405        const DOUBLE = 1 << 0;
2406
2407        /// The string is triple-quoted:
2408        /// it begins and ends with three consecutive quote characters.
2409        const TRIPLE_QUOTED = 1 << 1;
2410
2411        /// The string has a `u` or `U` prefix.
2412        /// While this prefix is a no-op at runtime,
2413        /// strings with this prefix can have no other prefixes set.
2414        const U_PREFIX = 1 << 2;
2415
2416        /// The string has a `b` or `B` prefix.
2417        /// This means that the string is a sequence of `int`s at runtime,
2418        /// rather than a sequence of `str`s.
2419        /// Strings with this flag can also be raw strings,
2420        /// but can have no other prefixes.
2421        const B_PREFIX = 1 << 3;
2422
2423        /// The string has a `f` or `F` prefix, meaning it is an f-string.
2424        /// F-strings can also be raw strings,
2425        /// but can have no other prefixes.
2426        const F_PREFIX = 1 << 4;
2427
2428        /// The string has a `t` or `T` prefix, meaning it is a t-string.
2429        /// T-strings can also be raw strings,
2430        /// but can have no other prefixes.
2431        const T_PREFIX = 1 << 5;
2432
2433        /// The string has an `r` prefix, meaning it is a raw string.
2434        /// F-strings and byte-strings can be raw,
2435        /// as can strings with no other prefixes.
2436        /// U-strings cannot be raw.
2437        const R_PREFIX_LOWER = 1 << 6;
2438
2439        /// The string has an `R` prefix, meaning it is a raw string.
2440        /// The casing of the `r`/`R` has no semantic significance at runtime;
2441        /// see https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#r-strings-and-r-strings
2442        /// for why we track the casing of the `r` prefix,
2443        /// but not for any other prefix
2444        const R_PREFIX_UPPER = 1 << 7;
2445
2446        /// String without matching closing quote(s).
2447        const UNCLOSED = 1 << 8;
2448    }
2449}
2450
2451#[derive(Clone, Copy, PartialEq, Eq, Hash)]
2452pub struct AnyStringFlags(AnyStringFlagsInner);
2453
2454impl AnyStringFlags {
2455    #[must_use]
2456    pub fn with_prefix(mut self, prefix: AnyStringPrefix) -> Self {
2457        self.0 |= match prefix {
2458            // regular strings
2459            AnyStringPrefix::Regular(StringLiteralPrefix::Empty) => AnyStringFlagsInner::empty(),
2460            AnyStringPrefix::Regular(StringLiteralPrefix::Unicode) => AnyStringFlagsInner::U_PREFIX,
2461            AnyStringPrefix::Regular(StringLiteralPrefix::Raw { uppercase: false }) => {
2462                AnyStringFlagsInner::R_PREFIX_LOWER
2463            }
2464            AnyStringPrefix::Regular(StringLiteralPrefix::Raw { uppercase: true }) => {
2465                AnyStringFlagsInner::R_PREFIX_UPPER
2466            }
2467
2468            // bytestrings
2469            AnyStringPrefix::Bytes(ByteStringPrefix::Regular) => AnyStringFlagsInner::B_PREFIX,
2470            AnyStringPrefix::Bytes(ByteStringPrefix::Raw { uppercase_r: false }) => {
2471                AnyStringFlagsInner::B_PREFIX.union(AnyStringFlagsInner::R_PREFIX_LOWER)
2472            }
2473            AnyStringPrefix::Bytes(ByteStringPrefix::Raw { uppercase_r: true }) => {
2474                AnyStringFlagsInner::B_PREFIX.union(AnyStringFlagsInner::R_PREFIX_UPPER)
2475            }
2476
2477            // f-strings
2478            AnyStringPrefix::Format(FStringPrefix::Regular) => AnyStringFlagsInner::F_PREFIX,
2479            AnyStringPrefix::Format(FStringPrefix::Raw { uppercase_r: false }) => {
2480                AnyStringFlagsInner::F_PREFIX.union(AnyStringFlagsInner::R_PREFIX_LOWER)
2481            }
2482            AnyStringPrefix::Format(FStringPrefix::Raw { uppercase_r: true }) => {
2483                AnyStringFlagsInner::F_PREFIX.union(AnyStringFlagsInner::R_PREFIX_UPPER)
2484            }
2485
2486            // t-strings
2487            AnyStringPrefix::Template(TStringPrefix::Regular) => AnyStringFlagsInner::T_PREFIX,
2488            AnyStringPrefix::Template(TStringPrefix::Raw { uppercase_r: false }) => {
2489                AnyStringFlagsInner::T_PREFIX.union(AnyStringFlagsInner::R_PREFIX_LOWER)
2490            }
2491            AnyStringPrefix::Template(TStringPrefix::Raw { uppercase_r: true }) => {
2492                AnyStringFlagsInner::T_PREFIX.union(AnyStringFlagsInner::R_PREFIX_UPPER)
2493            }
2494        };
2495        self
2496    }
2497
2498    pub fn new(prefix: AnyStringPrefix, quotes: Quote, triple_quotes: TripleQuotes) -> Self {
2499        Self(AnyStringFlagsInner::empty())
2500            .with_prefix(prefix)
2501            .with_quote_style(quotes)
2502            .with_triple_quotes(triple_quotes)
2503    }
2504
2505    /// Does the string have a `u` or `U` prefix?
2506    pub const fn is_u_string(self) -> bool {
2507        self.0.contains(AnyStringFlagsInner::U_PREFIX)
2508    }
2509
2510    /// Does the string have an `r` or `R` prefix?
2511    pub const fn is_raw_string(self) -> bool {
2512        self.0.intersects(
2513            AnyStringFlagsInner::R_PREFIX_LOWER.union(AnyStringFlagsInner::R_PREFIX_UPPER),
2514        )
2515    }
2516
2517    /// Does the string have an `f`,`F`,`t`, or `T` prefix?
2518    pub const fn is_interpolated_string(self) -> bool {
2519        self.0
2520            .intersects(AnyStringFlagsInner::F_PREFIX.union(AnyStringFlagsInner::T_PREFIX))
2521    }
2522
2523    /// Does the string have a `b` or `B` prefix?
2524    pub const fn is_byte_string(self) -> bool {
2525        self.0.contains(AnyStringFlagsInner::B_PREFIX)
2526    }
2527
2528    #[must_use]
2529    pub fn with_quote_style(mut self, quotes: Quote) -> Self {
2530        match quotes {
2531            Quote::Double => self.0 |= AnyStringFlagsInner::DOUBLE,
2532            Quote::Single => self.0 -= AnyStringFlagsInner::DOUBLE,
2533        }
2534        self
2535    }
2536
2537    #[must_use]
2538    pub fn with_triple_quotes(mut self, triple_quotes: TripleQuotes) -> Self {
2539        self.0
2540            .set(AnyStringFlagsInner::TRIPLE_QUOTED, triple_quotes.is_yes());
2541        self
2542    }
2543
2544    #[must_use]
2545    pub fn with_unclosed(mut self, unclosed: bool) -> Self {
2546        self.0.set(AnyStringFlagsInner::UNCLOSED, unclosed);
2547        self
2548    }
2549}
2550
2551impl StringFlags for AnyStringFlags {
2552    /// Does the string use single or double quotes in its opener and closer?
2553    fn quote_style(self) -> Quote {
2554        if self.0.contains(AnyStringFlagsInner::DOUBLE) {
2555            Quote::Double
2556        } else {
2557            Quote::Single
2558        }
2559    }
2560
2561    fn triple_quotes(self) -> TripleQuotes {
2562        if self.0.contains(AnyStringFlagsInner::TRIPLE_QUOTED) {
2563            TripleQuotes::Yes
2564        } else {
2565            TripleQuotes::No
2566        }
2567    }
2568
2569    fn prefix(self) -> AnyStringPrefix {
2570        let AnyStringFlags(flags) = self;
2571
2572        // f-strings
2573        if flags.contains(AnyStringFlagsInner::F_PREFIX) {
2574            if flags.contains(AnyStringFlagsInner::R_PREFIX_LOWER) {
2575                return AnyStringPrefix::Format(FStringPrefix::Raw { uppercase_r: false });
2576            }
2577            if flags.contains(AnyStringFlagsInner::R_PREFIX_UPPER) {
2578                return AnyStringPrefix::Format(FStringPrefix::Raw { uppercase_r: true });
2579            }
2580            return AnyStringPrefix::Format(FStringPrefix::Regular);
2581        }
2582
2583        // t-strings
2584        if flags.contains(AnyStringFlagsInner::T_PREFIX) {
2585            if flags.contains(AnyStringFlagsInner::R_PREFIX_LOWER) {
2586                return AnyStringPrefix::Template(TStringPrefix::Raw { uppercase_r: false });
2587            }
2588            if flags.contains(AnyStringFlagsInner::R_PREFIX_UPPER) {
2589                return AnyStringPrefix::Template(TStringPrefix::Raw { uppercase_r: true });
2590            }
2591            return AnyStringPrefix::Template(TStringPrefix::Regular);
2592        }
2593
2594        // bytestrings
2595        if flags.contains(AnyStringFlagsInner::B_PREFIX) {
2596            if flags.contains(AnyStringFlagsInner::R_PREFIX_LOWER) {
2597                return AnyStringPrefix::Bytes(ByteStringPrefix::Raw { uppercase_r: false });
2598            }
2599            if flags.contains(AnyStringFlagsInner::R_PREFIX_UPPER) {
2600                return AnyStringPrefix::Bytes(ByteStringPrefix::Raw { uppercase_r: true });
2601            }
2602            return AnyStringPrefix::Bytes(ByteStringPrefix::Regular);
2603        }
2604
2605        // all other strings
2606        if flags.contains(AnyStringFlagsInner::R_PREFIX_LOWER) {
2607            return AnyStringPrefix::Regular(StringLiteralPrefix::Raw { uppercase: false });
2608        }
2609        if flags.contains(AnyStringFlagsInner::R_PREFIX_UPPER) {
2610            return AnyStringPrefix::Regular(StringLiteralPrefix::Raw { uppercase: true });
2611        }
2612        if flags.contains(AnyStringFlagsInner::U_PREFIX) {
2613            return AnyStringPrefix::Regular(StringLiteralPrefix::Unicode);
2614        }
2615        AnyStringPrefix::Regular(StringLiteralPrefix::Empty)
2616    }
2617
2618    fn is_unclosed(self) -> bool {
2619        self.0.intersects(AnyStringFlagsInner::UNCLOSED)
2620    }
2621}
2622
2623impl fmt::Debug for AnyStringFlags {
2624    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2625        f.debug_struct("AnyStringFlags")
2626            .field("prefix", &self.prefix())
2627            .field("triple_quoted", &self.is_triple_quoted())
2628            .field("quote_style", &self.quote_style())
2629            .field("unclosed", &self.is_unclosed())
2630            .finish()
2631    }
2632}
2633
2634impl From<AnyStringFlags> for StringLiteralFlags {
2635    fn from(value: AnyStringFlags) -> StringLiteralFlags {
2636        let AnyStringPrefix::Regular(prefix) = value.prefix() else {
2637            unreachable!(
2638                "Should never attempt to convert {} into a regular string",
2639                value.prefix()
2640            )
2641        };
2642        StringLiteralFlags::empty()
2643            .with_quote_style(value.quote_style())
2644            .with_prefix(prefix)
2645            .with_triple_quotes(value.triple_quotes())
2646            .with_unclosed(value.is_unclosed())
2647    }
2648}
2649
2650impl From<StringLiteralFlags> for AnyStringFlags {
2651    fn from(value: StringLiteralFlags) -> Self {
2652        value.as_any_string_flags()
2653    }
2654}
2655
2656impl From<AnyStringFlags> for BytesLiteralFlags {
2657    fn from(value: AnyStringFlags) -> BytesLiteralFlags {
2658        let AnyStringPrefix::Bytes(bytestring_prefix) = value.prefix() else {
2659            unreachable!(
2660                "Should never attempt to convert {} into a bytestring",
2661                value.prefix()
2662            )
2663        };
2664        BytesLiteralFlags::empty()
2665            .with_quote_style(value.quote_style())
2666            .with_prefix(bytestring_prefix)
2667            .with_triple_quotes(value.triple_quotes())
2668            .with_unclosed(value.is_unclosed())
2669    }
2670}
2671
2672impl From<BytesLiteralFlags> for AnyStringFlags {
2673    fn from(value: BytesLiteralFlags) -> Self {
2674        value.as_any_string_flags()
2675    }
2676}
2677
2678impl From<AnyStringFlags> for FStringFlags {
2679    fn from(value: AnyStringFlags) -> FStringFlags {
2680        let AnyStringPrefix::Format(prefix) = value.prefix() else {
2681            unreachable!(
2682                "Should never attempt to convert {} into an f-string",
2683                value.prefix()
2684            )
2685        };
2686        FStringFlags::empty()
2687            .with_quote_style(value.quote_style())
2688            .with_prefix(prefix)
2689            .with_triple_quotes(value.triple_quotes())
2690            .with_unclosed(value.is_unclosed())
2691    }
2692}
2693
2694impl From<FStringFlags> for AnyStringFlags {
2695    fn from(value: FStringFlags) -> Self {
2696        value.as_any_string_flags()
2697    }
2698}
2699
2700impl From<AnyStringFlags> for TStringFlags {
2701    fn from(value: AnyStringFlags) -> TStringFlags {
2702        let AnyStringPrefix::Template(prefix) = value.prefix() else {
2703            unreachable!(
2704                "Should never attempt to convert {} into a t-string",
2705                value.prefix()
2706            )
2707        };
2708        TStringFlags::empty()
2709            .with_quote_style(value.quote_style())
2710            .with_prefix(prefix)
2711            .with_triple_quotes(value.triple_quotes())
2712            .with_unclosed(value.is_unclosed())
2713    }
2714}
2715
2716impl From<TStringFlags> for AnyStringFlags {
2717    fn from(value: TStringFlags) -> Self {
2718        value.as_any_string_flags()
2719    }
2720}
2721
2722#[derive(Clone, Debug, PartialEq, is_macro::Is)]
2723#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
2724pub enum Number {
2725    Int(int::Int),
2726    Float(f64),
2727    Complex { real: f64, imag: f64 },
2728}
2729
2730impl ExprName {
2731    pub fn id(&self) -> &Name {
2732        &self.id
2733    }
2734
2735    /// Returns `true` if this node represents an invalid name i.e., the `ctx` is [`Invalid`].
2736    ///
2737    /// [`Invalid`]: ExprContext::Invalid
2738    pub const fn is_invalid(&self) -> bool {
2739        matches!(self.ctx, ExprContext::Invalid)
2740    }
2741}
2742
2743impl ExprList {
2744    pub fn iter(&self) -> std::slice::Iter<'_, Expr> {
2745        self.elts.iter()
2746    }
2747
2748    pub fn len(&self) -> usize {
2749        self.elts.len()
2750    }
2751
2752    pub fn is_empty(&self) -> bool {
2753        self.elts.is_empty()
2754    }
2755}
2756
2757impl<'a> IntoIterator for &'a ExprList {
2758    type IntoIter = std::slice::Iter<'a, Expr>;
2759    type Item = &'a Expr;
2760
2761    fn into_iter(self) -> Self::IntoIter {
2762        self.iter()
2763    }
2764}
2765
2766impl ExprTuple {
2767    pub fn iter(&self) -> std::slice::Iter<'_, Expr> {
2768        self.elts.iter()
2769    }
2770
2771    pub fn len(&self) -> usize {
2772        self.elts.len()
2773    }
2774
2775    pub fn is_empty(&self) -> bool {
2776        self.elts.is_empty()
2777    }
2778}
2779
2780impl<'a> IntoIterator for &'a ExprTuple {
2781    type IntoIter = std::slice::Iter<'a, Expr>;
2782    type Item = &'a Expr;
2783
2784    fn into_iter(self) -> Self::IntoIter {
2785        self.iter()
2786    }
2787}
2788
2789/// See also [expr_context](https://docs.python.org/3/library/ast.html#ast.expr_context)
2790#[derive(Clone, Debug, PartialEq, is_macro::Is, Copy, Hash, Eq)]
2791#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
2792pub enum ExprContext {
2793    Load,
2794    Store,
2795    Del,
2796    Invalid,
2797}
2798
2799/// See also [boolop](https://docs.python.org/3/library/ast.html#ast.BoolOp)
2800#[derive(Clone, Debug, PartialEq, is_macro::Is, Copy, Hash, Eq)]
2801#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
2802pub enum BoolOp {
2803    And,
2804    Or,
2805}
2806
2807impl BoolOp {
2808    pub const fn as_str(&self) -> &'static str {
2809        match self {
2810            BoolOp::And => "and",
2811            BoolOp::Or => "or",
2812        }
2813    }
2814}
2815
2816impl fmt::Display for BoolOp {
2817    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2818        f.write_str(self.as_str())
2819    }
2820}
2821
2822/// See also [operator](https://docs.python.org/3/library/ast.html#ast.operator)
2823#[derive(Clone, Debug, PartialEq, is_macro::Is, Copy, Hash, Eq)]
2824#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
2825pub enum Operator {
2826    Add,
2827    Sub,
2828    Mult,
2829    MatMult,
2830    Div,
2831    Mod,
2832    Pow,
2833    LShift,
2834    RShift,
2835    BitOr,
2836    BitXor,
2837    BitAnd,
2838    FloorDiv,
2839}
2840
2841impl Operator {
2842    pub const fn as_str(&self) -> &'static str {
2843        match self {
2844            Operator::Add => "+",
2845            Operator::Sub => "-",
2846            Operator::Mult => "*",
2847            Operator::MatMult => "@",
2848            Operator::Div => "/",
2849            Operator::Mod => "%",
2850            Operator::Pow => "**",
2851            Operator::LShift => "<<",
2852            Operator::RShift => ">>",
2853            Operator::BitOr => "|",
2854            Operator::BitXor => "^",
2855            Operator::BitAnd => "&",
2856            Operator::FloorDiv => "//",
2857        }
2858    }
2859
2860    /// Returns the dunder method name for the operator.
2861    pub const fn dunder(self) -> &'static str {
2862        match self {
2863            Operator::Add => "__add__",
2864            Operator::Sub => "__sub__",
2865            Operator::Mult => "__mul__",
2866            Operator::MatMult => "__matmul__",
2867            Operator::Div => "__truediv__",
2868            Operator::Mod => "__mod__",
2869            Operator::Pow => "__pow__",
2870            Operator::LShift => "__lshift__",
2871            Operator::RShift => "__rshift__",
2872            Operator::BitOr => "__or__",
2873            Operator::BitXor => "__xor__",
2874            Operator::BitAnd => "__and__",
2875            Operator::FloorDiv => "__floordiv__",
2876        }
2877    }
2878
2879    /// Returns the in-place dunder method name for the operator.
2880    pub const fn in_place_dunder(self) -> &'static str {
2881        match self {
2882            Operator::Add => "__iadd__",
2883            Operator::Sub => "__isub__",
2884            Operator::Mult => "__imul__",
2885            Operator::MatMult => "__imatmul__",
2886            Operator::Div => "__itruediv__",
2887            Operator::Mod => "__imod__",
2888            Operator::Pow => "__ipow__",
2889            Operator::LShift => "__ilshift__",
2890            Operator::RShift => "__irshift__",
2891            Operator::BitOr => "__ior__",
2892            Operator::BitXor => "__ixor__",
2893            Operator::BitAnd => "__iand__",
2894            Operator::FloorDiv => "__ifloordiv__",
2895        }
2896    }
2897
2898    /// Returns the reflected dunder method name for the operator.
2899    pub const fn reflected_dunder(self) -> &'static str {
2900        match self {
2901            Operator::Add => "__radd__",
2902            Operator::Sub => "__rsub__",
2903            Operator::Mult => "__rmul__",
2904            Operator::MatMult => "__rmatmul__",
2905            Operator::Div => "__rtruediv__",
2906            Operator::Mod => "__rmod__",
2907            Operator::Pow => "__rpow__",
2908            Operator::LShift => "__rlshift__",
2909            Operator::RShift => "__rrshift__",
2910            Operator::BitOr => "__ror__",
2911            Operator::BitXor => "__rxor__",
2912            Operator::BitAnd => "__rand__",
2913            Operator::FloorDiv => "__rfloordiv__",
2914        }
2915    }
2916}
2917
2918impl fmt::Display for Operator {
2919    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2920        f.write_str(self.as_str())
2921    }
2922}
2923
2924/// See also [unaryop](https://docs.python.org/3/library/ast.html#ast.unaryop)
2925#[derive(Clone, Debug, PartialEq, is_macro::Is, Copy, Hash, Eq)]
2926#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
2927pub enum UnaryOp {
2928    Invert,
2929    Not,
2930    UAdd,
2931    USub,
2932}
2933
2934impl UnaryOp {
2935    pub const fn as_str(&self) -> &'static str {
2936        match self {
2937            UnaryOp::Invert => "~",
2938            UnaryOp::Not => "not",
2939            UnaryOp::UAdd => "+",
2940            UnaryOp::USub => "-",
2941        }
2942    }
2943}
2944
2945impl fmt::Display for UnaryOp {
2946    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2947        f.write_str(self.as_str())
2948    }
2949}
2950
2951/// See also [cmpop](https://docs.python.org/3/library/ast.html#ast.cmpop)
2952#[derive(Clone, Debug, PartialEq, is_macro::Is, Copy, Hash, Eq)]
2953#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
2954pub enum CmpOp {
2955    Eq,
2956    NotEq,
2957    Lt,
2958    LtE,
2959    Gt,
2960    GtE,
2961    Is,
2962    IsNot,
2963    In,
2964    NotIn,
2965}
2966
2967impl CmpOp {
2968    pub const fn as_str(&self) -> &'static str {
2969        match self {
2970            CmpOp::Eq => "==",
2971            CmpOp::NotEq => "!=",
2972            CmpOp::Lt => "<",
2973            CmpOp::LtE => "<=",
2974            CmpOp::Gt => ">",
2975            CmpOp::GtE => ">=",
2976            CmpOp::Is => "is",
2977            CmpOp::IsNot => "is not",
2978            CmpOp::In => "in",
2979            CmpOp::NotIn => "not in",
2980        }
2981    }
2982
2983    #[must_use]
2984    pub const fn negate(&self) -> Self {
2985        match self {
2986            CmpOp::Eq => CmpOp::NotEq,
2987            CmpOp::NotEq => CmpOp::Eq,
2988            CmpOp::Lt => CmpOp::GtE,
2989            CmpOp::LtE => CmpOp::Gt,
2990            CmpOp::Gt => CmpOp::LtE,
2991            CmpOp::GtE => CmpOp::Lt,
2992            CmpOp::Is => CmpOp::IsNot,
2993            CmpOp::IsNot => CmpOp::Is,
2994            CmpOp::In => CmpOp::NotIn,
2995            CmpOp::NotIn => CmpOp::In,
2996        }
2997    }
2998}
2999
3000impl fmt::Display for CmpOp {
3001    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3002        f.write_str(self.as_str())
3003    }
3004}
3005
3006/// See also [comprehension](https://docs.python.org/3/library/ast.html#ast.comprehension)
3007#[derive(Clone, Debug, PartialEq)]
3008#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3009pub struct Comprehension {
3010    pub range: TextRange,
3011    pub node_index: AtomicNodeIndex,
3012    pub target: Expr,
3013    pub iter: Expr,
3014    pub ifs: Vec<Expr>,
3015    pub is_async: bool,
3016    pub runtime_ifs: Option<Vec<Option<Expr>>>,
3017    pub runtime_is_async: Option<i32>,
3018}
3019
3020/// See also [ExceptHandler](https://docs.python.org/3/library/ast.html#ast.ExceptHandler)
3021#[derive(Clone, Debug, PartialEq)]
3022#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3023pub struct ExceptHandlerExceptHandler {
3024    pub range: TextRange,
3025    pub node_index: AtomicNodeIndex,
3026    pub type_: Option<Box<Expr>>,
3027    pub name: Option<Identifier>,
3028    pub body: Suite,
3029    pub runtime_body: Option<Vec<Option<Stmt>>>,
3030}
3031
3032/// See also [arg](https://docs.python.org/3/library/ast.html#ast.arg)
3033#[derive(Clone, Debug, PartialEq)]
3034#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3035pub struct Parameter {
3036    pub range: TextRange,
3037    pub node_index: AtomicNodeIndex,
3038    pub name: Identifier,
3039    pub annotation: Option<Box<Expr>>,
3040    pub runtime_type_comment: Option<Box<str>>,
3041    pub runtime_type_comment_bytes: Option<Vec<u8>>,
3042}
3043
3044impl Parameter {
3045    pub const fn name(&self) -> &Identifier {
3046        &self.name
3047    }
3048
3049    pub fn annotation(&self) -> Option<&Expr> {
3050        self.annotation.as_deref()
3051    }
3052}
3053
3054/// See also [keyword](https://docs.python.org/3/library/ast.html#ast.keyword)
3055#[derive(Clone, Debug, PartialEq)]
3056#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3057pub struct Keyword {
3058    pub range: TextRange,
3059    pub node_index: AtomicNodeIndex,
3060    pub arg: Option<Identifier>,
3061    pub value: Expr,
3062}
3063
3064/// See also [alias](https://docs.python.org/3/library/ast.html#ast.alias)
3065#[derive(Clone, Debug, PartialEq)]
3066#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3067pub struct Alias {
3068    pub range: TextRange,
3069    pub node_index: AtomicNodeIndex,
3070    pub name: Identifier,
3071    pub asname: Option<Identifier>,
3072}
3073
3074/// See also [withitem](https://docs.python.org/3/library/ast.html#ast.withitem)
3075#[derive(Clone, Debug, PartialEq)]
3076#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3077pub struct WithItem {
3078    pub range: TextRange,
3079    pub node_index: AtomicNodeIndex,
3080    pub context_expr: Expr,
3081    pub optional_vars: Option<Box<Expr>>,
3082}
3083
3084/// See also [match_case](https://docs.python.org/3/library/ast.html#ast.match_case)
3085#[derive(Clone, Debug, PartialEq)]
3086#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3087pub struct MatchCase {
3088    pub range: TextRange,
3089    pub node_index: AtomicNodeIndex,
3090    pub pattern: Pattern,
3091    pub guard: Option<Box<Expr>>,
3092    pub body: Suite,
3093    pub runtime_body: Option<Vec<Option<Stmt>>>,
3094}
3095
3096impl Pattern {
3097    /// Checks if the [`Pattern`] is an [irrefutable pattern].
3098    ///
3099    /// [irrefutable pattern]: https://peps.python.org/pep-0634/#irrefutable-case-blocks
3100    pub fn is_irrefutable(&self) -> bool {
3101        self.irrefutable_pattern().is_some()
3102    }
3103
3104    /// Return `Some(IrrefutablePattern)` if `self` is irrefutable or `None` otherwise.
3105    pub fn irrefutable_pattern(&self) -> Option<IrrefutablePattern> {
3106        match self {
3107            Pattern::MatchAs(PatternMatchAs {
3108                pattern,
3109                name,
3110                range,
3111                node_index,
3112            }) => match pattern {
3113                Some(pattern) => pattern.irrefutable_pattern(),
3114                None => match name {
3115                    Some(name) => Some(IrrefutablePattern {
3116                        kind: IrrefutablePatternKind::Name(name.id.clone()),
3117                        range: *range,
3118                        node_index: node_index.clone(),
3119                    }),
3120                    None => Some(IrrefutablePattern {
3121                        kind: IrrefutablePatternKind::Wildcard,
3122                        range: *range,
3123                        node_index: node_index.clone(),
3124                    }),
3125                },
3126            },
3127            Pattern::MatchOr(PatternMatchOr { patterns, .. }) => {
3128                patterns.iter().find_map(Pattern::irrefutable_pattern)
3129            }
3130            _ => None,
3131        }
3132    }
3133
3134    /// Checks if the [`Pattern`] is a [wildcard pattern].
3135    ///
3136    /// The following are wildcard patterns:
3137    /// ```python
3138    /// match subject:
3139    ///     case _ as x: ...
3140    ///     case _ | _: ...
3141    ///     case _: ...
3142    /// ```
3143    ///
3144    /// [wildcard pattern]: https://docs.python.org/3/reference/compound_stmts.html#wildcard-patterns
3145    pub fn is_wildcard(&self) -> bool {
3146        match self {
3147            Pattern::MatchAs(PatternMatchAs { pattern, .. }) => {
3148                pattern.as_deref().is_none_or(Pattern::is_wildcard)
3149            }
3150            Pattern::MatchOr(PatternMatchOr { patterns, .. }) => {
3151                patterns.iter().all(Pattern::is_wildcard)
3152            }
3153            _ => false,
3154        }
3155    }
3156}
3157
3158pub struct IrrefutablePattern {
3159    pub kind: IrrefutablePatternKind,
3160    pub range: TextRange,
3161    pub node_index: AtomicNodeIndex,
3162}
3163
3164#[derive(Debug, Clone, PartialEq, Eq, Hash)]
3165#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3166pub enum IrrefutablePatternKind {
3167    Name(Name),
3168    Wildcard,
3169}
3170
3171/// An AST node to represent the arguments to a [`crate::PatternMatchClass`], i.e., the
3172/// parenthesized contents in `case Point(1, x=0, y=0)`.
3173///
3174/// Like [`Arguments`], but for [`crate::PatternMatchClass`].
3175#[derive(Clone, Debug, PartialEq)]
3176#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3177pub struct PatternArguments {
3178    pub range: TextRange,
3179    pub node_index: AtomicNodeIndex,
3180    pub patterns: ThinVec<Pattern>,
3181    pub keywords: Vec<PatternKeyword>,
3182}
3183
3184/// An AST node to represent the keyword arguments to a [`crate::PatternMatchClass`], i.e., the
3185/// `x=0` and `y=0` in `case Point(x=0, y=0)`.
3186///
3187/// Like [`Keyword`], but for [`crate::PatternMatchClass`].
3188#[derive(Clone, Debug, PartialEq)]
3189#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3190pub struct PatternKeyword {
3191    pub range: TextRange,
3192    pub node_index: AtomicNodeIndex,
3193    pub attr: Identifier,
3194    pub pattern: Pattern,
3195}
3196
3197impl PatternArguments {
3198    /// Returns an iterator over the patterns and keywords in source order.
3199    pub fn iter_source_order(&self) -> PatternArgumentsSourceOrder<'_> {
3200        PatternArgumentsSourceOrder {
3201            patterns: &self.patterns,
3202            keywords: &self.keywords,
3203            next_pattern: 0,
3204            next_keyword: 0,
3205        }
3206    }
3207}
3208
3209/// The iterator returned by [`PatternArguments::iter_source_order`].
3210#[derive(Clone)]
3211pub struct PatternArgumentsSourceOrder<'a> {
3212    patterns: &'a [Pattern],
3213    keywords: &'a [PatternKeyword],
3214    next_pattern: usize,
3215    next_keyword: usize,
3216}
3217
3218/// An entry in the argument list of a class pattern.
3219#[derive(Copy, Clone, Debug, PartialEq)]
3220pub enum PatternOrKeyword<'a> {
3221    Pattern(&'a Pattern),
3222    Keyword(&'a PatternKeyword),
3223}
3224
3225impl<'a> Iterator for PatternArgumentsSourceOrder<'a> {
3226    type Item = PatternOrKeyword<'a>;
3227
3228    fn next(&mut self) -> Option<Self::Item> {
3229        let pattern = self.patterns.get(self.next_pattern);
3230        let keyword = self.keywords.get(self.next_keyword);
3231
3232        if let Some(pattern) = pattern
3233            && keyword.is_none_or(|keyword| pattern.start() <= keyword.start())
3234        {
3235            self.next_pattern += 1;
3236            Some(PatternOrKeyword::Pattern(pattern))
3237        } else if let Some(keyword) = keyword {
3238            self.next_keyword += 1;
3239            Some(PatternOrKeyword::Keyword(keyword))
3240        } else {
3241            None
3242        }
3243    }
3244}
3245
3246impl FusedIterator for PatternArgumentsSourceOrder<'_> {}
3247
3248impl TypeParam {
3249    pub const fn name(&self) -> &Identifier {
3250        match self {
3251            Self::TypeVar(x) => &x.name,
3252            Self::ParamSpec(x) => &x.name,
3253            Self::TypeVarTuple(x) => &x.name,
3254        }
3255    }
3256
3257    pub fn default(&self) -> Option<&Expr> {
3258        match self {
3259            Self::TypeVar(x) => x.default.as_deref(),
3260            Self::ParamSpec(x) => x.default.as_deref(),
3261            Self::TypeVarTuple(x) => x.default.as_deref(),
3262        }
3263    }
3264}
3265
3266/// See also [decorator](https://docs.python.org/3/library/ast.html#ast.decorator)
3267#[derive(Clone, Debug, PartialEq)]
3268#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3269pub struct Decorator {
3270    pub range: TextRange,
3271    pub node_index: AtomicNodeIndex,
3272    pub expression: Expr,
3273}
3274
3275/// Enumeration of the two kinds of parameter
3276#[derive(Debug, PartialEq, Clone, Copy)]
3277pub enum AnyParameterRef<'a> {
3278    /// Variadic parameters cannot have default values,
3279    /// e.g. both `*args` and `**kwargs` in the following function:
3280    ///
3281    /// ```python
3282    /// def foo(*args, **kwargs): pass
3283    /// ```
3284    Variadic(&'a Parameter),
3285
3286    /// Non-variadic parameters can have default values,
3287    /// though they won't necessarily always have them:
3288    ///
3289    /// ```python
3290    /// def bar(a=1, /, b=2, *, c=3): pass
3291    /// ```
3292    NonVariadic(&'a ParameterWithDefault),
3293}
3294
3295impl<'a> AnyParameterRef<'a> {
3296    pub const fn as_parameter(self) -> &'a Parameter {
3297        match self {
3298            Self::NonVariadic(param) => &param.parameter,
3299            Self::Variadic(param) => param,
3300        }
3301    }
3302
3303    pub const fn name(self) -> &'a Identifier {
3304        &self.as_parameter().name
3305    }
3306
3307    pub const fn is_variadic(self) -> bool {
3308        matches!(self, Self::Variadic(_))
3309    }
3310
3311    pub fn annotation(self) -> Option<&'a Expr> {
3312        self.as_parameter().annotation.as_deref()
3313    }
3314
3315    pub fn default(self) -> Option<&'a Expr> {
3316        match self {
3317            Self::NonVariadic(param) => param.default.as_deref(),
3318            Self::Variadic(_) => None,
3319        }
3320    }
3321}
3322
3323impl Ranged for AnyParameterRef<'_> {
3324    fn range(&self) -> TextRange {
3325        match self {
3326            Self::NonVariadic(param) => param.range,
3327            Self::Variadic(param) => param.range,
3328        }
3329    }
3330}
3331
3332/// An alternative type of AST `arguments`. This is ruff_python_parser-friendly and human-friendly definition of function arguments.
3333/// This form also has advantage to implement pre-order traverse.
3334///
3335/// `defaults` and `kw_defaults` fields are removed and the default values are placed under each [`ParameterWithDefault`] typed argument.
3336/// `vararg` and `kwarg` are still typed as `arg` because they never can have a default value.
3337///
3338/// The original Python-style AST type orders `kwonlyargs` fields by default existence; [Parameters] has location-ordered `kwonlyargs` fields.
3339///
3340/// NOTE: This type differs from the original Python AST. See: [arguments](https://docs.python.org/3/library/ast.html#ast.arguments).
3341
3342#[derive(Clone, Debug, PartialEq, Default)]
3343#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3344pub struct Parameters {
3345    pub range: TextRange,
3346    pub node_index: AtomicNodeIndex,
3347    pub posonlyargs: ThinVec<ParameterWithDefault>,
3348    pub args: ThinVec<ParameterWithDefault>,
3349    pub vararg: Option<Box<Parameter>>,
3350    pub kwonlyargs: ThinVec<ParameterWithDefault>,
3351    pub kwarg: Option<Box<Parameter>>,
3352    pub runtime_defaults: Option<Vec<Option<Expr>>>,
3353}
3354
3355impl Parameters {
3356    /// Returns an iterator over all non-variadic parameters included in this [`Parameters`] node.
3357    ///
3358    /// The variadic parameters (`.vararg` and `.kwarg`) can never have default values;
3359    /// non-variadic parameters sometimes will.
3360    pub fn iter_non_variadic_params(&self) -> impl Iterator<Item = &ParameterWithDefault> {
3361        self.posonlyargs
3362            .iter()
3363            .chain(&self.args)
3364            .chain(&self.kwonlyargs)
3365    }
3366
3367    /// Returns the [`ParameterWithDefault`] with the given name, or `None` if no such [`ParameterWithDefault`] exists.
3368    pub fn find(&self, name: &str) -> Option<&ParameterWithDefault> {
3369        self.iter_non_variadic_params()
3370            .find(|arg| arg.parameter.name.as_str() == name)
3371    }
3372
3373    /// Returns the index of the parameter with the given name
3374    pub fn index(&self, name: &str) -> Option<usize> {
3375        self.iter_non_variadic_params()
3376            .position(|arg| arg.parameter.name.as_str() == name)
3377    }
3378
3379    /// Returns an iterator over all parameters included in this [`Parameters`] node.
3380    pub fn iter(&self) -> ParametersIterator<'_> {
3381        ParametersIterator::new(self)
3382    }
3383
3384    /// Returns the total number of parameters included in this [`Parameters`] node.
3385    pub fn len(&self) -> usize {
3386        let Parameters {
3387            range: _,
3388            node_index: _,
3389            posonlyargs,
3390            args,
3391            vararg,
3392            kwonlyargs,
3393            kwarg,
3394            ..
3395        } = self;
3396        // Safety: a Python function can have an arbitrary number of parameters,
3397        // so theoretically this could be a number that wouldn't fit into a usize,
3398        // which would lead to a panic. A Python function with that many parameters
3399        // is extremely unlikely outside of generated code, however, and it's even
3400        // more unlikely that we'd find a function with that many parameters in a
3401        // source-code file <=4GB large (Ruff's maximum).
3402        posonlyargs
3403            .len()
3404            .checked_add(args.len())
3405            .and_then(|length| length.checked_add(usize::from(vararg.is_some())))
3406            .and_then(|length| length.checked_add(kwonlyargs.len()))
3407            .and_then(|length| length.checked_add(usize::from(kwarg.is_some())))
3408            .expect("Failed to fit the number of parameters into a usize")
3409    }
3410
3411    /// Returns `true` if a parameter with the given name is included in this [`Parameters`].
3412    pub fn includes(&self, name: &str) -> bool {
3413        self.iter().any(|param| param.name() == name)
3414    }
3415
3416    /// Returns `true` if the [`Parameters`] is empty.
3417    pub fn is_empty(&self) -> bool {
3418        self.posonlyargs.is_empty()
3419            && self.args.is_empty()
3420            && self.kwonlyargs.is_empty()
3421            && self.vararg.is_none()
3422            && self.kwarg.is_none()
3423    }
3424
3425    /// Returns an iterator over all parameters in source order.
3426    ///
3427    /// This differs from [`Parameters::iter`] which returns parameters
3428    /// in type-based order (positional-only, regular, variadic, keyword-only,
3429    /// keyword). For well-formed Python the two orderings are identical, but
3430    /// error recovery can produce ASTs where variadic parameters appear before
3431    /// non-variadic ones (e.g. `def foo(**kwargs, a):`).
3432    pub fn iter_source_order(&self) -> ParametersSourceOrderIterator<'_> {
3433        let mut variadics = [self.vararg.as_deref(), self.kwarg.as_deref()];
3434        variadics.sort_by_key(|param| param.map_or(TextSize::new(u32::MAX), Ranged::start));
3435
3436        ParametersSourceOrderIterator {
3437            next_non_variadic_peeked: None,
3438            posonlyargs: self.posonlyargs.iter(),
3439            args: self.args.iter(),
3440            kwonlyargs: self.kwonlyargs.iter(),
3441            variadics,
3442            next_variadic: 0,
3443        }
3444    }
3445}
3446
3447pub struct ParametersIterator<'a> {
3448    posonlyargs: Iter<'a, ParameterWithDefault>,
3449    args: Iter<'a, ParameterWithDefault>,
3450    vararg: Option<&'a Parameter>,
3451    kwonlyargs: Iter<'a, ParameterWithDefault>,
3452    kwarg: Option<&'a Parameter>,
3453}
3454
3455impl<'a> ParametersIterator<'a> {
3456    fn new(parameters: &'a Parameters) -> Self {
3457        let Parameters {
3458            range: _,
3459            node_index: _,
3460            posonlyargs,
3461            args,
3462            vararg,
3463            kwonlyargs,
3464            kwarg,
3465            ..
3466        } = parameters;
3467        Self {
3468            posonlyargs: posonlyargs.iter(),
3469            args: args.iter(),
3470            vararg: vararg.as_deref(),
3471            kwonlyargs: kwonlyargs.iter(),
3472            kwarg: kwarg.as_deref(),
3473        }
3474    }
3475}
3476
3477impl<'a> Iterator for ParametersIterator<'a> {
3478    type Item = AnyParameterRef<'a>;
3479
3480    fn next(&mut self) -> Option<Self::Item> {
3481        let ParametersIterator {
3482            posonlyargs,
3483            args,
3484            vararg,
3485            kwonlyargs,
3486            kwarg,
3487        } = self;
3488
3489        if let Some(param) = posonlyargs.next() {
3490            return Some(AnyParameterRef::NonVariadic(param));
3491        }
3492        if let Some(param) = args.next() {
3493            return Some(AnyParameterRef::NonVariadic(param));
3494        }
3495        if let Some(param) = vararg.take() {
3496            return Some(AnyParameterRef::Variadic(param));
3497        }
3498        if let Some(param) = kwonlyargs.next() {
3499            return Some(AnyParameterRef::NonVariadic(param));
3500        }
3501        kwarg.take().map(AnyParameterRef::Variadic)
3502    }
3503
3504    fn size_hint(&self) -> (usize, Option<usize>) {
3505        let ParametersIterator {
3506            posonlyargs,
3507            args,
3508            vararg,
3509            kwonlyargs,
3510            kwarg,
3511        } = self;
3512
3513        let posonlyargs_len = posonlyargs.len();
3514        let args_len = args.len();
3515        let vararg_len = usize::from(vararg.is_some());
3516        let kwonlyargs_len = kwonlyargs.len();
3517        let kwarg_len = usize::from(kwarg.is_some());
3518
3519        let lower = posonlyargs_len
3520            .saturating_add(args_len)
3521            .saturating_add(vararg_len)
3522            .saturating_add(kwonlyargs_len)
3523            .saturating_add(kwarg_len);
3524
3525        let upper = posonlyargs_len
3526            .checked_add(args_len)
3527            .and_then(|length| length.checked_add(vararg_len))
3528            .and_then(|length| length.checked_add(kwonlyargs_len))
3529            .and_then(|length| length.checked_add(kwarg_len));
3530
3531        (lower, upper)
3532    }
3533
3534    fn last(mut self) -> Option<Self::Item> {
3535        self.next_back()
3536    }
3537}
3538
3539impl DoubleEndedIterator for ParametersIterator<'_> {
3540    fn next_back(&mut self) -> Option<Self::Item> {
3541        let ParametersIterator {
3542            posonlyargs,
3543            args,
3544            vararg,
3545            kwonlyargs,
3546            kwarg,
3547        } = self;
3548
3549        if let Some(param) = kwarg.take() {
3550            return Some(AnyParameterRef::Variadic(param));
3551        }
3552        if let Some(param) = kwonlyargs.next_back() {
3553            return Some(AnyParameterRef::NonVariadic(param));
3554        }
3555        if let Some(param) = vararg.take() {
3556            return Some(AnyParameterRef::Variadic(param));
3557        }
3558        if let Some(param) = args.next_back() {
3559            return Some(AnyParameterRef::NonVariadic(param));
3560        }
3561        posonlyargs.next_back().map(AnyParameterRef::NonVariadic)
3562    }
3563}
3564
3565impl FusedIterator for ParametersIterator<'_> {}
3566
3567/// We rely on the same invariants outlined in the comment above `Parameters::len()`
3568/// in order to implement `ExactSizeIterator` here
3569impl ExactSizeIterator for ParametersIterator<'_> {}
3570
3571impl<'a> IntoIterator for &'a Parameters {
3572    type IntoIter = ParametersIterator<'a>;
3573    type Item = AnyParameterRef<'a>;
3574    fn into_iter(self) -> Self::IntoIter {
3575        self.iter()
3576    }
3577}
3578
3579impl<'a> IntoIterator for &'a Box<Parameters> {
3580    type IntoIter = ParametersIterator<'a>;
3581    type Item = AnyParameterRef<'a>;
3582    fn into_iter(self) -> Self::IntoIter {
3583        (&**self).into_iter()
3584    }
3585}
3586
3587/// The iterator returned by [`Parameters::iter_source_order`].
3588pub struct ParametersSourceOrderIterator<'a> {
3589    next_non_variadic_peeked: Option<&'a ParameterWithDefault>,
3590    posonlyargs: Iter<'a, ParameterWithDefault>,
3591    args: Iter<'a, ParameterWithDefault>,
3592    kwonlyargs: Iter<'a, ParameterWithDefault>,
3593    variadics: [Option<&'a Parameter>; 2],
3594    next_variadic: usize,
3595}
3596
3597impl<'a> ParametersSourceOrderIterator<'a> {
3598    /// Returns the next variadic parameter that appears before `before`, if any.
3599    fn next_variadic_before(&mut self, before: TextSize) -> Option<&'a Parameter> {
3600        let param = self.variadics.get(self.next_variadic).copied().flatten()?;
3601        if param.start() < before {
3602            self.next_variadic += 1;
3603            Some(param)
3604        } else {
3605            None
3606        }
3607    }
3608
3609    fn next_non_variadic(&mut self) -> Option<&'a ParameterWithDefault> {
3610        self.next_non_variadic_peeked
3611            .take()
3612            .or_else(|| self.posonlyargs.next())
3613            .or_else(|| self.args.next())
3614            .or_else(|| self.kwonlyargs.next())
3615    }
3616
3617    fn peek_next_non_variadic(&mut self) -> Option<&'a ParameterWithDefault> {
3618        let next = self.next_non_variadic()?;
3619        self.next_non_variadic_peeked = Some(next);
3620        Some(next)
3621    }
3622}
3623
3624impl<'a> Iterator for ParametersSourceOrderIterator<'a> {
3625    type Item = AnyParameterRef<'a>;
3626
3627    fn next(&mut self) -> Option<Self::Item> {
3628        // If there's a variadic parameter that comes before the next
3629        // non-variadic parameter, emit it first.
3630        let next_non_variadic_start = self
3631            .peek_next_non_variadic()
3632            .map_or(TextSize::new(u32::MAX), Ranged::start);
3633
3634        if let Some(variadic) = self.next_variadic_before(next_non_variadic_start) {
3635            return Some(AnyParameterRef::Variadic(variadic));
3636        }
3637
3638        if let Some(non_variadic) = self.next_non_variadic() {
3639            return Some(AnyParameterRef::NonVariadic(non_variadic));
3640        }
3641
3642        // Drain remaining variadics.
3643        self.next_variadic_before(TextSize::new(u32::MAX))
3644            .map(AnyParameterRef::Variadic)
3645    }
3646}
3647
3648impl FusedIterator for ParametersSourceOrderIterator<'_> {}
3649
3650/// An alternative type of AST `arg`. This is used for each function argument that might have a default value.
3651/// Used by `Arguments` original type.
3652///
3653/// NOTE: This type is different from original Python AST.
3654#[derive(Clone, Debug, PartialEq)]
3655#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3656pub struct ParameterWithDefault {
3657    pub range: TextRange,
3658    pub node_index: AtomicNodeIndex,
3659    pub parameter: Parameter,
3660    pub default: Option<Box<Expr>>,
3661}
3662
3663impl ParameterWithDefault {
3664    pub fn default(&self) -> Option<&Expr> {
3665        self.default.as_deref()
3666    }
3667
3668    pub const fn name(&self) -> &Identifier {
3669        self.parameter.name()
3670    }
3671
3672    pub fn annotation(&self) -> Option<&Expr> {
3673        self.parameter.annotation()
3674    }
3675
3676    /// Return `true` if the parameter name uses the pre-PEP-570 convention
3677    /// (specified in PEP 484) to indicate to a type checker that it should be treated
3678    /// as positional-only.
3679    pub fn uses_pep_484_positional_only_convention(&self) -> bool {
3680        let name = self.name();
3681        name.starts_with("__") && !name.ends_with("__")
3682    }
3683}
3684
3685/// An AST node used to represent the arguments passed to a function call or class definition.
3686///
3687/// For example, given:
3688/// ```python
3689/// foo(1, 2, 3, bar=4, baz=5)
3690/// ```
3691/// The `Arguments` node would span from the left to right parentheses (inclusive), and contain
3692/// the arguments and keyword arguments in the order they appear in the source code.
3693///
3694/// Similarly, given:
3695/// ```python
3696/// class Foo(Bar, baz=1, qux=2):
3697///     pass
3698/// ```
3699/// The `Arguments` node would again span from the left to right parentheses (inclusive), and
3700/// contain the `Bar` argument and the `baz` and `qux` keyword arguments in the order they
3701/// appear in the source code.
3702///
3703/// In the context of a class definition, the Python-style AST refers to the arguments as `bases`,
3704/// as they represent the "explicitly specified base classes", while the keyword arguments are
3705/// typically used for `metaclass`, with any additional arguments being passed to the `metaclass`.
3706
3707#[derive(Clone, Debug, PartialEq)]
3708#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3709pub struct Arguments {
3710    pub range: TextRange,
3711    pub node_index: AtomicNodeIndex,
3712    pub args: ThinVec<Expr>,
3713    pub keywords: ThinVec<Keyword>,
3714    pub runtime_args: Option<Vec<Option<Expr>>>,
3715    pub runtime_bases: Option<Vec<Option<Expr>>>,
3716}
3717
3718/// An entry in the argument list of a function call.
3719#[derive(Copy, Clone, Debug, PartialEq)]
3720pub enum ArgOrKeyword<'a> {
3721    Arg(&'a Expr),
3722    Keyword(&'a Keyword),
3723}
3724
3725impl<'a> ArgOrKeyword<'a> {
3726    pub const fn value(self) -> &'a Expr {
3727        match self {
3728            ArgOrKeyword::Arg(argument) => argument,
3729            ArgOrKeyword::Keyword(keyword) => &keyword.value,
3730        }
3731    }
3732
3733    pub const fn is_variadic(self) -> bool {
3734        match self {
3735            ArgOrKeyword::Arg(expr) => expr.is_starred_expr(),
3736            ArgOrKeyword::Keyword(keyword) => keyword.arg.is_none(),
3737        }
3738    }
3739
3740    pub const fn as_variadic(self) -> Option<&'a Keyword> {
3741        match self {
3742            ArgOrKeyword::Keyword(keyword) if keyword.arg.is_none() => Some(keyword),
3743            _ => None,
3744        }
3745    }
3746
3747    pub const fn as_keyword(self) -> Option<&'a Keyword> {
3748        match self {
3749            ArgOrKeyword::Keyword(keyword) => Some(keyword),
3750            ArgOrKeyword::Arg(_) => None,
3751        }
3752    }
3753}
3754
3755impl<'a> From<&'a Expr> for ArgOrKeyword<'a> {
3756    fn from(arg: &'a Expr) -> Self {
3757        Self::Arg(arg)
3758    }
3759}
3760
3761impl<'a> From<&'a Keyword> for ArgOrKeyword<'a> {
3762    fn from(keyword: &'a Keyword) -> Self {
3763        Self::Keyword(keyword)
3764    }
3765}
3766
3767impl Ranged for ArgOrKeyword<'_> {
3768    fn range(&self) -> TextRange {
3769        match self {
3770            Self::Arg(arg) => arg.range(),
3771            Self::Keyword(keyword) => keyword.range(),
3772        }
3773    }
3774}
3775
3776impl Arguments {
3777    /// Return the number of positional and keyword arguments.
3778    pub fn len(&self) -> usize {
3779        self.args.len() + self.keywords.len()
3780    }
3781
3782    /// Return `true` if there are no positional or keyword arguments.
3783    pub fn is_empty(&self) -> bool {
3784        self.len() == 0
3785    }
3786
3787    /// Return the [`Keyword`] with the given name, or `None` if no such [`Keyword`] exists.
3788    pub fn find_keyword(&self, keyword_name: &str) -> Option<&Keyword> {
3789        self.keywords.iter().find(|keyword| {
3790            let Keyword { arg, .. } = keyword;
3791            arg.as_ref().is_some_and(|arg| arg == keyword_name)
3792        })
3793    }
3794
3795    /// Return the positional argument at the given index, or `None` if no such argument exists.
3796    pub fn find_positional(&self, position: usize) -> Option<&Expr> {
3797        self.args
3798            .iter()
3799            .take_while(|expr| !expr.is_starred_expr())
3800            .nth(position)
3801    }
3802
3803    /// Return the value for the argument with the given name or at the given position, or `None` if no such
3804    /// argument exists. Used to retrieve argument values that can be provided _either_ as keyword or
3805    /// positional arguments.
3806    pub fn find_argument_value(&self, name: &str, position: usize) -> Option<&Expr> {
3807        self.find_argument(name, position).map(ArgOrKeyword::value)
3808    }
3809
3810    /// Return the argument with the given name or at the given position, or `None` if no such
3811    /// argument exists. Used to retrieve arguments that can be provided _either_ as keyword or
3812    /// positional arguments.
3813    pub fn find_argument(&self, name: &str, position: usize) -> Option<ArgOrKeyword<'_>> {
3814        self.find_keyword(name)
3815            .map(ArgOrKeyword::from)
3816            .or_else(|| self.find_positional(position).map(ArgOrKeyword::from))
3817    }
3818
3819    /// Iterates over the positional and keyword arguments in the order of declaration.
3820    ///
3821    /// Positional arguments are generally before keyword arguments, but star arguments are an
3822    /// exception:
3823    /// ```python
3824    /// class A(*args, a=2, *args2, **kwargs):
3825    ///     pass
3826    ///
3827    /// f(*args, a=2, *args2, **kwargs)
3828    /// ```
3829    /// where `*args` and `args2` are `args` while `a=1` and `kwargs` are `keywords`.
3830    ///
3831    /// If you would just chain `args` and `keywords` the call would get reordered which we don't
3832    /// want. This function instead "merge sorts" them into the correct order.
3833    ///
3834    /// Note that the order of evaluation is always first `args`, then `keywords`:
3835    /// ```python
3836    /// def f(*args, **kwargs):
3837    ///     pass
3838    ///
3839    /// def g(x):
3840    ///     print(x)
3841    ///     return x
3842    ///
3843    ///
3844    /// f(*g([1]), a=g(2), *g([3]), **g({"4": 5}))
3845    /// ```
3846    /// Output:
3847    /// ```text
3848    /// [1]
3849    /// [3]
3850    /// 2
3851    /// {'4': 5}
3852    /// ```
3853    pub fn iter_source_order(&self) -> ArgumentsSourceOrder<'_> {
3854        ArgumentsSourceOrder {
3855            args: &self.args,
3856            keywords: &self.keywords,
3857            next_arg: 0,
3858            next_keyword: 0,
3859        }
3860    }
3861
3862    pub fn inner_range(&self) -> TextRange {
3863        TextRange::new(self.l_paren_range().end(), self.r_paren_range().start())
3864    }
3865
3866    pub fn l_paren_range(&self) -> TextRange {
3867        TextRange::at(self.start(), '('.text_len())
3868    }
3869
3870    pub fn r_paren_range(&self) -> TextRange {
3871        TextRange::new(self.end() - ')'.text_len(), self.end())
3872    }
3873}
3874
3875/// The iterator returned by [`Arguments::iter_source_order`].
3876#[derive(Clone)]
3877pub struct ArgumentsSourceOrder<'a> {
3878    args: &'a [Expr],
3879    keywords: &'a [Keyword],
3880    next_arg: usize,
3881    next_keyword: usize,
3882}
3883
3884impl<'a> Iterator for ArgumentsSourceOrder<'a> {
3885    type Item = ArgOrKeyword<'a>;
3886
3887    fn next(&mut self) -> Option<Self::Item> {
3888        let arg = self.args.get(self.next_arg);
3889        let keyword = self.keywords.get(self.next_keyword);
3890
3891        if let Some(arg) = arg
3892            && keyword.is_none_or(|keyword| arg.start() <= keyword.start())
3893        {
3894            self.next_arg += 1;
3895            Some(ArgOrKeyword::Arg(arg))
3896        } else if let Some(keyword) = keyword {
3897            self.next_keyword += 1;
3898            Some(ArgOrKeyword::Keyword(keyword))
3899        } else {
3900            None
3901        }
3902    }
3903}
3904
3905impl FusedIterator for ArgumentsSourceOrder<'_> {}
3906
3907/// An AST node used to represent a sequence of type parameters.
3908///
3909/// For example, given:
3910/// ```python
3911/// class C[T, U, V]: ...
3912/// ```
3913/// The `TypeParams` node would span from the left to right brackets (inclusive), and contain
3914/// the `T`, `U`, and `V` type parameters in the order they appear in the source code.
3915
3916#[derive(Clone, Debug, PartialEq)]
3917#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3918pub struct TypeParams {
3919    pub range: TextRange,
3920    pub node_index: AtomicNodeIndex,
3921    pub type_params: Vec<TypeParam>,
3922    pub runtime_type_params: Option<Vec<Option<TypeParam>>>,
3923}
3924
3925impl Deref for TypeParams {
3926    type Target = [TypeParam];
3927
3928    fn deref(&self) -> &Self::Target {
3929        &self.type_params
3930    }
3931}
3932
3933impl<'a> IntoIterator for &'a TypeParams {
3934    type Item = &'a TypeParam;
3935    type IntoIter = std::slice::Iter<'a, TypeParam>;
3936
3937    fn into_iter(self) -> Self::IntoIter {
3938        self.type_params.iter()
3939    }
3940}
3941
3942/// A suite represents a sequence of [`Stmt`].
3943///
3944/// See: <https://docs.python.org/3/reference/compound_stmts.html#grammar-token-python-grammar-suite>
3945pub type Suite = ThinVec<Stmt>;
3946
3947pub type DecoratorList = ThinVec<Decorator>;
3948
3949pub type Patterns = ThinVec<Pattern>;
3950
3951pub type PatternKeys = ThinVec<Expr>;
3952
3953pub type ParameterWithDefaults = ThinVec<ParameterWithDefault>;
3954
3955/// The kind of escape command as defined in [IPython Syntax] in the IPython codebase.
3956///
3957/// [IPython Syntax]: https://github.com/ipython/ipython/blob/635815e8f1ded5b764d66cacc80bbe25e9e2587f/IPython/core/inputtransformer2.py#L335-L343
3958#[derive(PartialEq, Eq, Debug, Clone, Hash, Copy)]
3959#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
3960pub enum IpyEscapeKind {
3961    /// Send line to underlying system shell (`!`).
3962    Shell,
3963    /// Send line to system shell and capture output (`!!`).
3964    ShCap,
3965    /// Show help on object (`?`).
3966    Help,
3967    /// Show help on object, with extra verbosity (`??`).
3968    Help2,
3969    /// Call magic function (`%`).
3970    Magic,
3971    /// Call cell magic function (`%%`).
3972    Magic2,
3973    /// Call first argument with rest of line as arguments after splitting on whitespace
3974    /// and quote each as string (`,`).
3975    Quote,
3976    /// Call first argument with rest of line as an argument quoted as a single string (`;`).
3977    Quote2,
3978    /// Call first argument with rest of line as arguments (`/`).
3979    Paren,
3980}
3981
3982impl TryFrom<char> for IpyEscapeKind {
3983    type Error = String;
3984
3985    fn try_from(ch: char) -> Result<Self, Self::Error> {
3986        match ch {
3987            '!' => Ok(IpyEscapeKind::Shell),
3988            '?' => Ok(IpyEscapeKind::Help),
3989            '%' => Ok(IpyEscapeKind::Magic),
3990            ',' => Ok(IpyEscapeKind::Quote),
3991            ';' => Ok(IpyEscapeKind::Quote2),
3992            '/' => Ok(IpyEscapeKind::Paren),
3993            _ => Err(format!("Unexpected magic escape: {ch}")),
3994        }
3995    }
3996}
3997
3998impl TryFrom<[char; 2]> for IpyEscapeKind {
3999    type Error = String;
4000
4001    fn try_from(ch: [char; 2]) -> Result<Self, Self::Error> {
4002        match ch {
4003            ['!', '!'] => Ok(IpyEscapeKind::ShCap),
4004            ['?', '?'] => Ok(IpyEscapeKind::Help2),
4005            ['%', '%'] => Ok(IpyEscapeKind::Magic2),
4006            [c1, c2] => Err(format!("Unexpected magic escape: {c1}{c2}")),
4007        }
4008    }
4009}
4010
4011impl fmt::Display for IpyEscapeKind {
4012    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
4013        f.write_str(self.as_str())
4014    }
4015}
4016
4017impl IpyEscapeKind {
4018    /// Returns `true` if the escape kind is help i.e., `?` or `??`.
4019    pub const fn is_help(self) -> bool {
4020        matches!(self, IpyEscapeKind::Help | IpyEscapeKind::Help2)
4021    }
4022
4023    /// Returns `true` if the escape kind is magic i.e., `%` or `%%`.
4024    pub const fn is_magic(self) -> bool {
4025        matches!(self, IpyEscapeKind::Magic | IpyEscapeKind::Magic2)
4026    }
4027
4028    pub fn as_str(self) -> &'static str {
4029        match self {
4030            IpyEscapeKind::Shell => "!",
4031            IpyEscapeKind::ShCap => "!!",
4032            IpyEscapeKind::Help => "?",
4033            IpyEscapeKind::Help2 => "??",
4034            IpyEscapeKind::Magic => "%",
4035            IpyEscapeKind::Magic2 => "%%",
4036            IpyEscapeKind::Quote => ",",
4037            IpyEscapeKind::Quote2 => ";",
4038            IpyEscapeKind::Paren => "/",
4039        }
4040    }
4041}
4042
4043/// An `Identifier` with an empty `id` is invalid.
4044///
4045/// For example, in the following code `id` will be empty.
4046/// ```python
4047/// def 1():
4048///     ...
4049/// ```
4050#[derive(Clone, Debug, PartialEq, Eq, Hash)]
4051#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
4052pub struct Identifier {
4053    pub id: Name,
4054    pub range: TextRange,
4055    pub node_index: AtomicNodeIndex,
4056}
4057
4058impl Identifier {
4059    #[inline]
4060    pub fn new(id: impl Into<Name>, range: TextRange) -> Self {
4061        Self {
4062            id: id.into(),
4063            node_index: AtomicNodeIndex::NONE,
4064            range,
4065        }
4066    }
4067
4068    pub fn id(&self) -> &Name {
4069        &self.id
4070    }
4071
4072    pub fn is_valid(&self) -> bool {
4073        !self.id.is_empty()
4074    }
4075}
4076
4077impl Identifier {
4078    #[inline]
4079    pub fn as_str(&self) -> &str {
4080        self.id.as_str()
4081    }
4082}
4083
4084impl PartialEq<str> for Identifier {
4085    #[inline]
4086    fn eq(&self, other: &str) -> bool {
4087        self.id == other
4088    }
4089}
4090
4091impl PartialEq<String> for Identifier {
4092    #[inline]
4093    fn eq(&self, other: &String) -> bool {
4094        self.id == other
4095    }
4096}
4097
4098impl std::ops::Deref for Identifier {
4099    type Target = str;
4100    #[inline]
4101    fn deref(&self) -> &Self::Target {
4102        self.id.as_str()
4103    }
4104}
4105
4106impl AsRef<str> for Identifier {
4107    #[inline]
4108    fn as_ref(&self) -> &str {
4109        self.id.as_str()
4110    }
4111}
4112
4113impl std::fmt::Display for Identifier {
4114    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4115        std::fmt::Display::fmt(&self.id, f)
4116    }
4117}
4118
4119impl From<Identifier> for Name {
4120    #[inline]
4121    fn from(identifier: Identifier) -> Name {
4122        identifier.id
4123    }
4124}
4125
4126#[derive(Clone, Copy, Debug, Hash, PartialEq)]
4127#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
4128pub enum Singleton {
4129    None,
4130    True,
4131    False,
4132}
4133
4134impl From<bool> for Singleton {
4135    fn from(value: bool) -> Self {
4136        if value {
4137            Singleton::True
4138        } else {
4139            Singleton::False
4140        }
4141    }
4142}
4143
4144#[cfg(test)]
4145mod tests {
4146    use ruff_text_size::{Ranged, TextRange, TextSize};
4147
4148    use crate::generated::*;
4149    use crate::{
4150        Arguments, AtomicNodeIndex, FString, FStringFlags, FStringPart, FStringPartMut,
4151        FStringValue, Mod, Parameters, StringLiteral, StringLiteralFlags,
4152    };
4153
4154    #[test]
4155    fn f_string_parts() {
4156        let fstring = |start| FString {
4157            range: TextRange::at(TextSize::new(start), TextSize::new(3)),
4158            node_index: AtomicNodeIndex::NONE,
4159            elements: Vec::new().into(),
4160            flags: FStringFlags::empty(),
4161        };
4162        let single = FStringValue::single(fstring(0));
4163        let concatenated = FStringValue::concatenated(vec![
4164            FStringPart::FString(fstring(0)),
4165            FStringPart::Literal(StringLiteral {
4166                range: TextRange::at(TextSize::new(3), TextSize::new(3)),
4167                node_index: AtomicNodeIndex::NONE,
4168                value: "x".into(),
4169                flags: StringLiteralFlags::empty(),
4170            }),
4171            FStringPart::FString(fstring(6)),
4172        ]);
4173
4174        for (mut value, starts) in [(single, vec![0]), (concatenated, vec![0, 3, 6])] {
4175            let mut parts = value.iter();
4176            let mut expected = starts.iter();
4177            while expected.len() > 0 {
4178                assert_eq!(parts.len(), expected.len());
4179                assert_eq!(
4180                    parts.next().map(|part| part.start().to_u32()),
4181                    expected.next().copied()
4182                );
4183                assert_eq!(
4184                    parts.next_back().map(|part| part.start().to_u32()),
4185                    expected.next_back().copied()
4186                );
4187            }
4188            assert_eq!(parts.len(), 0);
4189            assert_eq!(parts.next(), None);
4190            assert_eq!(parts.next_back(), None);
4191
4192            let mut parts = value.iter_mut();
4193            for (part, start) in parts.by_ref().rev().zip(starts.iter().rev()) {
4194                let range = match part {
4195                    FStringPartMut::Literal(literal) => &mut literal.range,
4196                    FStringPartMut::FString(fstring) => &mut fstring.range,
4197                };
4198                assert_eq!(range.start().to_u32(), *start);
4199                *range = TextRange::at(TextSize::new(start + 10), TextSize::new(3));
4200            }
4201            assert_eq!(parts.len(), 0);
4202            assert!(parts.next().is_none());
4203            assert!(parts.next_back().is_none());
4204            assert_eq!(
4205                value
4206                    .iter()
4207                    .map(|part| part.start().to_u32())
4208                    .collect::<Vec<_>>(),
4209                starts.iter().map(|start| start + 10).collect::<Vec<_>>(),
4210            );
4211        }
4212    }
4213
4214    #[test]
4215    #[cfg(target_pointer_width = "64")]
4216    fn size() {
4217        assert_eq!(std::mem::size_of::<Stmt>(), 176);
4218        assert_eq!(std::mem::size_of::<StmtFunctionDef>(), 176);
4219        assert_eq!(std::mem::size_of::<StmtClassDef>(), 128);
4220        assert_eq!(std::mem::size_of::<StmtTry>(), 160);
4221        assert_eq!(std::mem::size_of::<Mod>(), 48);
4222        assert_eq!(std::mem::size_of::<Pattern>(), 144);
4223        assert_eq!(std::mem::size_of::<Parameters>(), 80);
4224        assert_eq!(std::mem::size_of::<Arguments>(), 80);
4225        assert_eq!(std::mem::size_of::<Expr>(), 104);
4226        assert_eq!(std::mem::size_of::<ExprAttribute>(), 56);
4227        assert_eq!(std::mem::size_of::<ExprAwait>(), 24);
4228        assert_eq!(std::mem::size_of::<ExprBinOp>(), 32);
4229        assert_eq!(std::mem::size_of::<ExprBoolOp>(), 64);
4230        assert_eq!(std::mem::size_of::<ExprBooleanLiteral>(), 16);
4231        assert_eq!(std::mem::size_of::<ExprBytesLiteral>(), 48);
4232        assert_eq!(std::mem::size_of::<ExprCall>(), 96);
4233        assert_eq!(std::mem::size_of::<ExprCompare>(), 72);
4234        assert_eq!(std::mem::size_of::<ExprDict>(), 64);
4235        assert_eq!(std::mem::size_of::<ExprDictComp>(), 48);
4236        assert_eq!(std::mem::size_of::<ExprEllipsisLiteral>(), 12);
4237        assert_eq!(std::mem::size_of::<ExprFString>(), 96);
4238        assert_eq!(std::mem::size_of::<ExprGenerator>(), 48);
4239        assert_eq!(std::mem::size_of::<ExprIf>(), 40);
4240        assert_eq!(std::mem::size_of::<ExprIpyEscapeCommand>(), 32);
4241        assert_eq!(std::mem::size_of::<ExprLambda>(), 32);
4242        assert_eq!(std::mem::size_of::<ExprList>(), 64);
4243        assert_eq!(std::mem::size_of::<ExprListComp>(), 48);
4244        assert_eq!(std::mem::size_of::<ExprName>(), 32);
4245        assert_eq!(std::mem::size_of::<ExprNamed>(), 32);
4246        assert_eq!(std::mem::size_of::<ExprNoneLiteral>(), 12);
4247        assert_eq!(std::mem::size_of::<ExprNumberLiteral>(), 40);
4248        assert_eq!(std::mem::size_of::<ExprSet>(), 64);
4249        assert_eq!(std::mem::size_of::<ExprSetComp>(), 48);
4250        assert_eq!(std::mem::size_of::<ExprSlice>(), 40);
4251        assert_eq!(std::mem::size_of::<ExprStarred>(), 24);
4252        assert_eq!(std::mem::size_of::<ExprStringLiteral>(), 48);
4253        assert_eq!(std::mem::size_of::<ExprSubscript>(), 32);
4254        assert_eq!(std::mem::size_of::<ExprTString>(), 96);
4255        assert_eq!(std::mem::size_of::<ExprTuple>(), 64);
4256        assert_eq!(std::mem::size_of::<ExprUnaryOp>(), 24);
4257        assert_eq!(std::mem::size_of::<ExprYield>(), 24);
4258        assert_eq!(std::mem::size_of::<ExprYieldFrom>(), 24);
4259    }
4260}