From 03179b55ebe7e916f1722e18e8f0b87c01616d1f Mon Sep 17 00:00:00 2001
From: Zsolt Dollenstein <zsol.zsol@gmail.com>
Date: Sun, 1 Oct 2023 20:58:40 +0100
Subject: [PATCH] Parse arbitrarily nested f-strings (#1026)

---
 native/libcst/src/tokenizer/core/mod.rs       | 23 ++++++-------------
 native/libcst/src/tokenizer/tests.rs          | 22 ++++++++++++++++++
 native/libcst/tests/fixtures/super_strings.py |  9 +++++++-
 3 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/native/libcst/src/tokenizer/core/mod.rs b/native/libcst/src/tokenizer/core/mod.rs
index 2365eaa36..067c1cf9e 100644
--- a/native/libcst/src/tokenizer/core/mod.rs
+++ b/native/libcst/src/tokenizer/core/mod.rs
@@ -907,8 +907,8 @@ impl<'t> TokState<'t> {
                     return Err(TokError::UnterminatedString);
                 }
                 (ch @ Some('\''), _) | (ch @ Some('"'), _) => {
-                    // see if this actually terminates something in fstring_stack
-                    for node in self.fstring_stack.iter() {
+                    // see if this actually terminates the most recent fstring
+                    if let Some(node) = self.fstring_stack.last() {
                         if ch == Some(node.quote_char.into()) {
                             match node.quote_size {
                                 StringQuoteSize::Single => {
@@ -1004,27 +1004,18 @@ impl<'t> TokState<'t> {
 
     fn maybe_consume_fstring_end(&mut self) -> Option<TokType> {
         let ch = self.text_pos.peek();
-        let mut match_idx = None;
-        for (idx, node) in self.fstring_stack.iter().enumerate() {
+        if let Some(node) = self.fstring_stack.last() {
             if ch == Some(node.quote_char.into()) {
                 if node.quote_size == StringQuoteSize::Triple {
-                    if self.text_pos.consume(node.quote_char.triple_str()) {
-                        match_idx = Some(idx);
-                        break;
-                    }
+                    self.text_pos.consume(node.quote_char.triple_str());
                 } else {
                     self.text_pos.next(); // already matched
-                    match_idx = Some(idx);
-                    break;
                 }
+                self.fstring_stack.pop();
+                return Some(TokType::FStringEnd);
             }
         }
-        if let Some(match_idx) = match_idx {
-            self.fstring_stack.truncate(match_idx);
-            Some(TokType::FStringEnd)
-        } else {
-            None
-        }
+        None
     }
 }
 
diff --git a/native/libcst/src/tokenizer/tests.rs b/native/libcst/src/tokenizer/tests.rs
index a24b977b2..af79971df 100644
--- a/native/libcst/src/tokenizer/tests.rs
+++ b/native/libcst/src/tokenizer/tests.rs
@@ -853,3 +853,25 @@ fn test_nested_f_string_specs() {
         ])
     )
 }
+
+#[test]
+fn test_nested_f_strings() {
+    let config = TokConfig {
+        split_fstring: true,
+        ..default_config()
+    };
+    assert_eq!(
+        tokenize_all("f'{f'{2}'}'", &config),
+        Ok(vec![
+            (TokType::FStringStart, "f'"),
+            (TokType::Op, "{"),
+            (TokType::FStringStart, "f'"),
+            (TokType::Op, "{"),
+            (TokType::Number, "2"),
+            (TokType::Op, "}"),
+            (TokType::FStringEnd, "'"),
+            (TokType::Op, "}"),
+            (TokType::FStringEnd, "'")
+        ])
+    )
+}
diff --git a/native/libcst/tests/fixtures/super_strings.py b/native/libcst/tests/fixtures/super_strings.py
index 824572793..c1471a459 100644
--- a/native/libcst/tests/fixtures/super_strings.py
+++ b/native/libcst/tests/fixtures/super_strings.py
@@ -31,4 +31,11 @@
 f"regexp_like(path, '.*\{file_type}$')"
 f"\lfoo"
 
-f"{_:{_:}{a}}"
\ No newline at end of file
+f"{_:{_:}{a}}"
+
+f"foo {f"bar {x}"} baz"
+f'some words {a+b:.3f} more words {c+d=} final words'
+f"{'':*^{1:{1}}}"
+f"{'':*^{1:{1:{1}}}}"
+f"{f"{f"{f"{f"{f"{1+1}"}"}"}"}"}"
+