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