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