From 29fb86e3c831fd6e72283c0342f85e2b2fa7479c Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Tue, 24 Oct 2023 10:48:55 +0530
Subject: [PATCH] Allow unterminated f-strings in the indexer

---
 .../rules/implicit.rs                         | 33 ++++++++++++-------
 .../rules/invalid_escape_sequence.rs          | 24 +++++---------
 .../ruff_python_index/src/fstring_ranges.rs   |  4 ++-
 3 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/crates/ruff_linter/src/rules/flake8_implicit_str_concat/rules/implicit.rs b/crates/ruff_linter/src/rules/flake8_implicit_str_concat/rules/implicit.rs
index 14fff96dfb5dc..5e91f0ccead47 100644
--- a/crates/ruff_linter/src/rules/flake8_implicit_str_concat/rules/implicit.rs
+++ b/crates/ruff_linter/src/rules/flake8_implicit_str_concat/rules/implicit.rs
@@ -110,18 +110,27 @@ pub(crate) fn implicit(
     {
         let (a_range, b_range) = match (a_tok, b_tok) {
             (Tok::String { .. }, Tok::String { .. }) => (*a_range, *b_range),
-            (Tok::String { .. }, Tok::FStringStart) => (
-                *a_range,
-                indexer.fstring_ranges().innermost(b_range.start()).unwrap(),
-            ),
-            (Tok::FStringEnd, Tok::String { .. }) => (
-                indexer.fstring_ranges().innermost(a_range.start()).unwrap(),
-                *b_range,
-            ),
-            (Tok::FStringEnd, Tok::FStringStart) => (
-                indexer.fstring_ranges().innermost(a_range.start()).unwrap(),
-                indexer.fstring_ranges().innermost(b_range.start()).unwrap(),
-            ),
+            (Tok::String { .. }, Tok::FStringStart) => {
+                match indexer.fstring_ranges().innermost(b_range.start()) {
+                    Some(b_range) => (*a_range, b_range),
+                    None => continue,
+                }
+            }
+            (Tok::FStringEnd, Tok::String { .. }) => {
+                match indexer.fstring_ranges().innermost(a_range.start()) {
+                    Some(a_range) => (a_range, *b_range),
+                    None => continue,
+                }
+            }
+            (Tok::FStringEnd, Tok::FStringStart) => {
+                match (
+                    indexer.fstring_ranges().innermost(a_range.start()),
+                    indexer.fstring_ranges().innermost(b_range.start()),
+                ) {
+                    (Some(a_range), Some(b_range)) => (a_range, b_range),
+                    _ => continue,
+                }
+            }
             _ => continue,
         };
 
diff --git a/crates/ruff_linter/src/rules/pycodestyle/rules/invalid_escape_sequence.rs b/crates/ruff_linter/src/rules/pycodestyle/rules/invalid_escape_sequence.rs
index 4a333bdcd5206..d2d0f71c11460 100644
--- a/crates/ruff_linter/src/rules/pycodestyle/rules/invalid_escape_sequence.rs
+++ b/crates/ruff_linter/src/rules/pycodestyle/rules/invalid_escape_sequence.rs
@@ -50,18 +50,21 @@ pub(crate) fn invalid_escape_sequence(
     token: &Tok,
     token_range: TextRange,
 ) {
-    let token_source_code = match token {
+    let (token_source_code, string_start_location) = match token {
         Tok::FStringMiddle { value, is_raw } => {
             if *is_raw {
                 return;
             }
-            value.as_str()
+            let Some(range) = indexer.fstring_ranges().innermost(token_range.start()) else {
+                return;
+            };
+            (value.as_str(), range.start())
         }
         Tok::String { kind, .. } => {
             if kind.is_raw() {
                 return;
             }
-            locator.slice(token_range)
+            (locator.slice(token_range), token_range.start())
         }
         _ => return,
     };
@@ -166,25 +169,14 @@ pub(crate) fn invalid_escape_sequence(
             )));
         }
     } else {
-        let tok_start = if token.is_f_string_middle() {
-            // SAFETY: If this is a `FStringMiddle` token, then the indexer
-            // must have the f-string range.
-            indexer
-                .fstring_ranges()
-                .innermost(token_range.start())
-                .unwrap()
-                .start()
-        } else {
-            token_range.start()
-        };
         // Turn into raw string.
         for diagnostic in &mut invalid_escape_sequence {
             // If necessary, add a space between any leading keyword (`return`, `yield`,
             // `assert`, etc.) and the string. For example, `return"foo"` is valid, but
             // `returnr"foo"` is not.
             diagnostic.set_fix(Fix::safe_edit(Edit::insertion(
-                pad_start("r".to_string(), tok_start, locator),
-                tok_start,
+                pad_start("r".to_string(), string_start_location, locator),
+                string_start_location,
             )));
         }
     }
diff --git a/crates/ruff_python_index/src/fstring_ranges.rs b/crates/ruff_python_index/src/fstring_ranges.rs
index b95c9b665083b..bdc31258bb234 100644
--- a/crates/ruff_python_index/src/fstring_ranges.rs
+++ b/crates/ruff_python_index/src/fstring_ranges.rs
@@ -5,8 +5,11 @@ use ruff_text_size::{TextRange, TextSize};
 
 /// Stores the ranges of all f-strings in a file sorted by [`TextRange::start`].
 /// There can be multiple overlapping ranges for nested f-strings.
+///
+/// Note that the ranges for all unterminated f-strings are not stored.
 #[derive(Debug)]
 pub struct FStringRanges {
+    // Mapping from the f-string start location to its range.
     raw: BTreeMap<TextSize, TextRange>,
 }
 
@@ -89,7 +92,6 @@ impl FStringRangesBuilder {
     }
 
     pub(crate) fn finish(self) -> FStringRanges {
-        debug_assert!(self.start_locations.is_empty());
         FStringRanges { raw: self.raw }
     }
 }