ruff_python_ast/
expression.rs

1use std::iter::FusedIterator;
2
3use ruff_text_size::{Ranged, TextRange};
4
5use crate::{
6    self as ast, AnyNodeRef, AnyStringFlags, Expr, ExprBytesLiteral, ExprFString, ExprRef,
7    ExprStringLiteral, ExprTString, StringFlags,
8};
9
10impl<'a> From<&'a Box<Expr>> for ExprRef<'a> {
11    fn from(value: &'a Box<Expr>) -> Self {
12        ExprRef::from(value.as_ref())
13    }
14}
15
16/// Unowned pendant to all the literal variants of [`ast::Expr`] that stores a
17/// reference instead of an owned value.
18#[derive(Copy, Clone, Debug, PartialEq, is_macro::Is)]
19pub enum LiteralExpressionRef<'a> {
20    StringLiteral(&'a ast::ExprStringLiteral),
21    BytesLiteral(&'a ast::ExprBytesLiteral),
22    NumberLiteral(&'a ast::ExprNumberLiteral),
23    BooleanLiteral(&'a ast::ExprBooleanLiteral),
24    NoneLiteral(&'a ast::ExprNoneLiteral),
25    EllipsisLiteral(&'a ast::ExprEllipsisLiteral),
26}
27
28impl Ranged for LiteralExpressionRef<'_> {
29    fn range(&self) -> TextRange {
30        match self {
31            LiteralExpressionRef::StringLiteral(expression) => expression.range(),
32            LiteralExpressionRef::BytesLiteral(expression) => expression.range(),
33            LiteralExpressionRef::NumberLiteral(expression) => expression.range(),
34            LiteralExpressionRef::BooleanLiteral(expression) => expression.range(),
35            LiteralExpressionRef::NoneLiteral(expression) => expression.range(),
36            LiteralExpressionRef::EllipsisLiteral(expression) => expression.range(),
37        }
38    }
39}
40
41impl<'a> From<LiteralExpressionRef<'a>> for AnyNodeRef<'a> {
42    fn from(value: LiteralExpressionRef<'a>) -> Self {
43        match value {
44            LiteralExpressionRef::StringLiteral(expression) => {
45                AnyNodeRef::ExprStringLiteral(expression)
46            }
47            LiteralExpressionRef::BytesLiteral(expression) => {
48                AnyNodeRef::ExprBytesLiteral(expression)
49            }
50            LiteralExpressionRef::NumberLiteral(expression) => {
51                AnyNodeRef::ExprNumberLiteral(expression)
52            }
53            LiteralExpressionRef::BooleanLiteral(expression) => {
54                AnyNodeRef::ExprBooleanLiteral(expression)
55            }
56            LiteralExpressionRef::NoneLiteral(expression) => {
57                AnyNodeRef::ExprNoneLiteral(expression)
58            }
59            LiteralExpressionRef::EllipsisLiteral(expression) => {
60                AnyNodeRef::ExprEllipsisLiteral(expression)
61            }
62        }
63    }
64}
65
66impl LiteralExpressionRef<'_> {
67    /// Returns `true` if the literal is either a string or bytes literal that
68    /// is implicitly concatenated.
69    pub fn is_implicit_concatenated(&self) -> bool {
70        match self {
71            LiteralExpressionRef::StringLiteral(expression) => {
72                expression.value.is_implicit_concatenated()
73            }
74            LiteralExpressionRef::BytesLiteral(expression) => {
75                expression.value.is_implicit_concatenated()
76            }
77            _ => false,
78        }
79    }
80}
81
82/// An enum that holds a reference to a string-like expression from the AST. This includes string
83/// literals, bytes literals, f-strings, and t-strings.
84#[derive(Copy, Clone, Debug, PartialEq)]
85pub enum StringLike<'a> {
86    String(&'a ast::ExprStringLiteral),
87    Bytes(&'a ast::ExprBytesLiteral),
88    FString(&'a ast::ExprFString),
89    TString(&'a ast::ExprTString),
90}
91
92impl<'a> StringLike<'a> {
93    pub const fn is_interpolated_string(self) -> bool {
94        matches!(self, Self::TString(_) | Self::FString(_))
95    }
96
97    /// Returns an iterator over the [`StringLikePart`] contained in this string-like expression.
98    pub fn parts(&self) -> StringLikePartIter<'a> {
99        match self {
100            StringLike::String(expr) => StringLikePartIter::String(expr.value.iter()),
101            StringLike::Bytes(expr) => StringLikePartIter::Bytes(expr.value.iter()),
102            StringLike::FString(expr) => StringLikePartIter::FString(expr.value.iter()),
103            StringLike::TString(expr) => StringLikePartIter::TString(expr.value.iter()),
104        }
105    }
106
107    /// Returns `true` if the string is implicitly concatenated.
108    pub fn is_implicit_concatenated(self) -> bool {
109        match self {
110            Self::String(ExprStringLiteral { value, .. }) => value.is_implicit_concatenated(),
111            Self::Bytes(ExprBytesLiteral { value, .. }) => value.is_implicit_concatenated(),
112            Self::FString(ExprFString { value, .. }) => value.is_implicit_concatenated(),
113            Self::TString(ExprTString { value, .. }) => value.is_implicit_concatenated(),
114        }
115    }
116
117    pub const fn as_expression_ref(self) -> ExprRef<'a> {
118        match self {
119            StringLike::String(expr) => ExprRef::StringLiteral(expr),
120            StringLike::Bytes(expr) => ExprRef::BytesLiteral(expr),
121            StringLike::FString(expr) => ExprRef::FString(expr),
122            StringLike::TString(expr) => ExprRef::TString(expr),
123        }
124    }
125}
126
127impl<'a> From<&'a ast::ExprStringLiteral> for StringLike<'a> {
128    fn from(value: &'a ast::ExprStringLiteral) -> Self {
129        StringLike::String(value)
130    }
131}
132
133impl<'a> From<&'a ast::ExprBytesLiteral> for StringLike<'a> {
134    fn from(value: &'a ast::ExprBytesLiteral) -> Self {
135        StringLike::Bytes(value)
136    }
137}
138
139impl<'a> From<&'a ast::ExprFString> for StringLike<'a> {
140    fn from(value: &'a ast::ExprFString) -> Self {
141        StringLike::FString(value)
142    }
143}
144
145impl<'a> From<&'a ast::ExprTString> for StringLike<'a> {
146    fn from(value: &'a ast::ExprTString) -> Self {
147        StringLike::TString(value)
148    }
149}
150
151impl<'a> From<&StringLike<'a>> for ExprRef<'a> {
152    fn from(value: &StringLike<'a>) -> Self {
153        match value {
154            StringLike::String(expr) => ExprRef::StringLiteral(expr),
155            StringLike::Bytes(expr) => ExprRef::BytesLiteral(expr),
156            StringLike::FString(expr) => ExprRef::FString(expr),
157            StringLike::TString(expr) => ExprRef::TString(expr),
158        }
159    }
160}
161
162impl<'a> From<StringLike<'a>> for AnyNodeRef<'a> {
163    fn from(value: StringLike<'a>) -> Self {
164        AnyNodeRef::from(&value)
165    }
166}
167
168impl<'a> From<&StringLike<'a>> for AnyNodeRef<'a> {
169    fn from(value: &StringLike<'a>) -> Self {
170        match value {
171            StringLike::String(expr) => AnyNodeRef::ExprStringLiteral(expr),
172            StringLike::Bytes(expr) => AnyNodeRef::ExprBytesLiteral(expr),
173            StringLike::FString(expr) => AnyNodeRef::ExprFString(expr),
174            StringLike::TString(expr) => AnyNodeRef::ExprTString(expr),
175        }
176    }
177}
178
179impl<'a> TryFrom<&'a Expr> for StringLike<'a> {
180    type Error = ();
181
182    fn try_from(value: &'a Expr) -> Result<Self, Self::Error> {
183        match value {
184            Expr::StringLiteral(value) => Ok(Self::String(value)),
185            Expr::BytesLiteral(value) => Ok(Self::Bytes(value)),
186            Expr::FString(value) => Ok(Self::FString(value)),
187            Expr::TString(value) => Ok(Self::TString(value)),
188            _ => Err(()),
189        }
190    }
191}
192
193impl<'a> TryFrom<AnyNodeRef<'a>> for StringLike<'a> {
194    type Error = ();
195
196    fn try_from(value: AnyNodeRef<'a>) -> Result<Self, Self::Error> {
197        match value {
198            AnyNodeRef::ExprStringLiteral(value) => Ok(Self::String(value)),
199            AnyNodeRef::ExprBytesLiteral(value) => Ok(Self::Bytes(value)),
200            AnyNodeRef::ExprFString(value) => Ok(Self::FString(value)),
201            AnyNodeRef::ExprTString(value) => Ok(Self::TString(value)),
202            _ => Err(()),
203        }
204    }
205}
206
207impl Ranged for StringLike<'_> {
208    fn range(&self) -> TextRange {
209        match self {
210            StringLike::String(literal) => literal.range(),
211            StringLike::Bytes(literal) => literal.range(),
212            StringLike::FString(literal) => literal.range(),
213            StringLike::TString(literal) => literal.range(),
214        }
215    }
216}
217
218/// An enum that holds a reference to an individual part of a string-like expression.
219#[derive(Copy, Clone, Debug, PartialEq)]
220pub enum StringLikePart<'a> {
221    String(&'a ast::StringLiteral),
222    Bytes(&'a ast::BytesLiteral),
223    FString(&'a ast::FString),
224    TString(&'a ast::TString),
225}
226
227impl<'a> StringLikePart<'a> {
228    /// Returns the [`AnyStringFlags`] for the current string-like part.
229    pub fn flags(&self) -> AnyStringFlags {
230        match self {
231            StringLikePart::String(string) => AnyStringFlags::from(string.flags),
232            StringLikePart::Bytes(bytes) => AnyStringFlags::from(bytes.flags),
233            StringLikePart::FString(f_string) => AnyStringFlags::from(f_string.flags),
234            StringLikePart::TString(t_string) => AnyStringFlags::from(t_string.flags),
235        }
236    }
237
238    /// Returns the range of the string's content in the source (minus prefix and quotes).
239    pub fn content_range(self) -> TextRange {
240        let kind = self.flags();
241        TextRange::new(
242            self.start() + kind.opener_len(),
243            self.end() - kind.closer_len(),
244        )
245    }
246
247    pub const fn is_string_literal(self) -> bool {
248        matches!(self, Self::String(_))
249    }
250
251    pub const fn as_string_literal(self) -> Option<&'a ast::StringLiteral> {
252        match self {
253            StringLikePart::String(value) => Some(value),
254            _ => None,
255        }
256    }
257
258    pub const fn is_interpolated_string(self) -> bool {
259        matches!(self, Self::FString(_) | Self::TString(_))
260    }
261}
262
263impl<'a> From<&'a ast::StringLiteral> for StringLikePart<'a> {
264    fn from(value: &'a ast::StringLiteral) -> Self {
265        StringLikePart::String(value)
266    }
267}
268
269impl<'a> From<&'a ast::BytesLiteral> for StringLikePart<'a> {
270    fn from(value: &'a ast::BytesLiteral) -> Self {
271        StringLikePart::Bytes(value)
272    }
273}
274
275impl<'a> From<&'a ast::FString> for StringLikePart<'a> {
276    fn from(value: &'a ast::FString) -> Self {
277        StringLikePart::FString(value)
278    }
279}
280
281impl<'a> From<&'a ast::TString> for StringLikePart<'a> {
282    fn from(value: &'a ast::TString) -> Self {
283        StringLikePart::TString(value)
284    }
285}
286
287impl<'a> From<&StringLikePart<'a>> for AnyNodeRef<'a> {
288    fn from(value: &StringLikePart<'a>) -> Self {
289        AnyNodeRef::from(*value)
290    }
291}
292
293impl<'a> From<StringLikePart<'a>> for AnyNodeRef<'a> {
294    fn from(value: StringLikePart<'a>) -> Self {
295        match value {
296            StringLikePart::String(part) => AnyNodeRef::StringLiteral(part),
297            StringLikePart::Bytes(part) => AnyNodeRef::BytesLiteral(part),
298            StringLikePart::FString(part) => AnyNodeRef::FString(part),
299            StringLikePart::TString(part) => AnyNodeRef::TString(part),
300        }
301    }
302}
303
304impl Ranged for StringLikePart<'_> {
305    fn range(&self) -> TextRange {
306        match self {
307            StringLikePart::String(part) => part.range(),
308            StringLikePart::Bytes(part) => part.range(),
309            StringLikePart::FString(part) => part.range(),
310            StringLikePart::TString(part) => part.range(),
311        }
312    }
313}
314
315/// An iterator over all the [`StringLikePart`] of a string-like expression.
316///
317/// This is created by the [`StringLike::parts`] method.
318#[derive(Clone)]
319pub enum StringLikePartIter<'a> {
320    String(std::slice::Iter<'a, ast::StringLiteral>),
321    Bytes(std::slice::Iter<'a, ast::BytesLiteral>),
322    FString(std::slice::Iter<'a, ast::FStringPart>),
323    TString(std::slice::Iter<'a, ast::TString>),
324}
325
326impl<'a> Iterator for StringLikePartIter<'a> {
327    type Item = StringLikePart<'a>;
328
329    fn next(&mut self) -> Option<Self::Item> {
330        let part = match self {
331            StringLikePartIter::String(inner) => StringLikePart::String(inner.next()?),
332            StringLikePartIter::Bytes(inner) => StringLikePart::Bytes(inner.next()?),
333            StringLikePartIter::FString(inner) => {
334                let part = inner.next()?;
335                match part {
336                    ast::FStringPart::Literal(string_literal) => {
337                        StringLikePart::String(string_literal)
338                    }
339                    ast::FStringPart::FString(f_string) => StringLikePart::FString(f_string),
340                }
341            }
342            StringLikePartIter::TString(inner) => StringLikePart::TString(inner.next()?),
343        };
344
345        Some(part)
346    }
347
348    fn size_hint(&self) -> (usize, Option<usize>) {
349        match self {
350            StringLikePartIter::String(inner) => inner.size_hint(),
351            StringLikePartIter::Bytes(inner) => inner.size_hint(),
352            StringLikePartIter::FString(inner) => inner.size_hint(),
353            StringLikePartIter::TString(inner) => inner.size_hint(),
354        }
355    }
356}
357
358impl DoubleEndedIterator for StringLikePartIter<'_> {
359    fn next_back(&mut self) -> Option<Self::Item> {
360        let part = match self {
361            StringLikePartIter::String(inner) => StringLikePart::String(inner.next_back()?),
362            StringLikePartIter::Bytes(inner) => StringLikePart::Bytes(inner.next_back()?),
363            StringLikePartIter::FString(inner) => {
364                let part = inner.next_back()?;
365                match part {
366                    ast::FStringPart::Literal(string_literal) => {
367                        StringLikePart::String(string_literal)
368                    }
369                    ast::FStringPart::FString(f_string) => StringLikePart::FString(f_string),
370                }
371            }
372            StringLikePartIter::TString(inner) => StringLikePart::TString(inner.next_back()?),
373        };
374
375        Some(part)
376    }
377}
378
379impl FusedIterator for StringLikePartIter<'_> {}
380impl ExactSizeIterator for StringLikePartIter<'_> {}
ruff_python_ast/expression.rs

ruff_python_ast/
expression.rs
