Skip to content

Commit

Permalink
Add methods to iter over f-string elements (#10309)
Browse files Browse the repository at this point in the history
## Summary

This PR adds methods on `FString` to iterate over the two different kind
of elements it can have - literals and expressions. This is similar to
the methods we have on `ExprFString`.

---------

Co-authored-by: Alex Waygood <[email protected]>
  • Loading branch information
dhruvmanila and AlexWaygood authored Mar 13, 2024
1 parent 93d582d commit 32d6f84
Show file tree
Hide file tree
Showing 4 changed files with 19 additions and 13 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -168,11 +168,7 @@ fn should_be_fstring(

for f_string in value.f_strings() {
let mut has_name = false;
for element in f_string
.elements
.iter()
.filter_map(|element| element.as_expression())
{
for element in f_string.expressions() {
if let ast::Expr::Name(ast::ExprName { id, .. }) = element.expression.as_ref() {
if arg_names.contains(id.as_str()) {
return false;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -100,11 +100,7 @@ fn contains_message(expr: &Expr) -> bool {
}
}
ast::FStringPart::FString(f_string) => {
for literal in f_string
.elements
.iter()
.filter_map(|element| element.as_literal())
{
for literal in f_string.literals() {
if literal.chars().any(char::is_whitespace) {
return true;
}
Expand Down
16 changes: 16 additions & 0 deletions crates/ruff_python_ast/src/nodes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1248,6 +1248,22 @@ pub struct FString {
pub flags: FStringFlags,
}

impl FString {
/// Returns an iterator over all the [`FStringLiteralElement`] nodes contained in this f-string.
pub fn literals(&self) -> impl Iterator<Item = &FStringLiteralElement> {
self.elements
.iter()
.filter_map(|element| element.as_literal())
}

/// Returns an iterator over all the [`FStringExpressionElement`] nodes contained in this f-string.
pub fn expressions(&self) -> impl Iterator<Item = &FStringExpressionElement> {
self.elements
.iter()
.filter_map(|element| element.as_expression())
}
}

impl Ranged for FString {
fn range(&self) -> TextRange {
self.range
Expand Down
4 changes: 1 addition & 3 deletions crates/ruff_python_formatter/src/other/f_string.rs
Original file line number Diff line number Diff line change
Expand Up @@ -138,9 +138,7 @@ impl FStringLayout {
//
// Reference: https://prettier.io/docs/en/next/rationale.html#template-literals
if f_string
.elements
.iter()
.filter_map(|element| element.as_expression())
.expressions()
.any(|expr| memchr::memchr2(b'\n', b'\r', locator.slice(expr).as_bytes()).is_some())
{
Self::Multiline
Expand Down

0 comments on commit 32d6f84

Please sign in to comment.