Add methods to iter over f-string elements (#10309)

## 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 <Alex.Waygood@Gmail.com>
This commit is contained in:
Dhruv Manilawala 2024-03-13 14:16:55 +05:30 committed by GitHub
parent 93d582d734
commit 32d6f84e3d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 19 additions and 13 deletions

View file

@ -168,11 +168,7 @@ fn should_be_fstring(
for f_string in value.f_strings() { for f_string in value.f_strings() {
let mut has_name = false; let mut has_name = false;
for element in f_string for element in f_string.expressions() {
.elements
.iter()
.filter_map(|element| element.as_expression())
{
if let ast::Expr::Name(ast::ExprName { id, .. }) = element.expression.as_ref() { if let ast::Expr::Name(ast::ExprName { id, .. }) = element.expression.as_ref() {
if arg_names.contains(id.as_str()) { if arg_names.contains(id.as_str()) {
return false; return false;

View file

@ -100,11 +100,7 @@ fn contains_message(expr: &Expr) -> bool {
} }
} }
ast::FStringPart::FString(f_string) => { ast::FStringPart::FString(f_string) => {
for literal in f_string for literal in f_string.literals() {
.elements
.iter()
.filter_map(|element| element.as_literal())
{
if literal.chars().any(char::is_whitespace) { if literal.chars().any(char::is_whitespace) {
return true; return true;
} }

View file

@ -1248,6 +1248,22 @@ pub struct FString {
pub flags: FStringFlags, 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 { impl Ranged for FString {
fn range(&self) -> TextRange { fn range(&self) -> TextRange {
self.range self.range

View file

@ -138,9 +138,7 @@ impl FStringLayout {
// //
// Reference: https://prettier.io/docs/en/next/rationale.html#template-literals // Reference: https://prettier.io/docs/en/next/rationale.html#template-literals
if f_string if f_string
.elements .expressions()
.iter()
.filter_map(|element| element.as_expression())
.any(|expr| memchr::memchr2(b'\n', b'\r', locator.slice(expr).as_bytes()).is_some()) .any(|expr| memchr::memchr2(b'\n', b'\r', locator.slice(expr).as_bytes()).is_some())
{ {
Self::Multiline Self::Multiline