PY-31442 Spellchecker properly handles f-strings

It also now correctly processes escape sequences in
mixed up raw and normal string elements (previously
it stopped extracting escapes altogether once there
was at least one raw string element among those
belonging to a string literal expression).

python/src/com/jetbrains/python/spellchecker/PythonSpellcheckerStrategy.java

@@ -15,10 +15,7 @@
  */
 package com.jetbrains.python.spellchecker;
 
-import com.intellij.lang.ASTNode;
-import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.TextRange;
-import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
 import com.intellij.spellchecker.inspections.PlainTextSplitter;
 import com.intellij.spellchecker.inspections.Splitter;
@@ -29,14 +26,15 @@ import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.inspections.PyStringFormatParser;
 import com.jetbrains.python.psi.PyBinaryExpression;
+import com.jetbrains.python.psi.PyFormattedStringElement;
+import com.jetbrains.python.psi.PyStringElement;
 import com.jetbrains.python.psi.PyStringLiteralExpression;
-import com.jetbrains.python.psi.PyStringLiteralUtil;
+import com.jetbrains.python.psi.impl.PyStringLiteralDecoder;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.Collections;
 import java.util.List;
 
-import static com.jetbrains.python.psi.PyUtil.StringNodeInfo;
-
 /**
  * @author yole
  */
@@ -45,25 +43,29 @@ public class PythonSpellcheckerStrategy extends SpellcheckingStrategy {
     @Override
     public void tokenize(@NotNull PyStringLiteralExpression element, TokenConsumer consumer) {
       final Splitter splitter = PlainTextSplitter.getInstance();
-      final List<ASTNode> strNodes = element.getStringNodes();
-      final List<String> prefixes = ContainerUtil.mapNotNull(strNodes, n -> StringUtil.nullize(new StringNodeInfo(n).getPrefix()));
-      
-      if (element.textContains('\\') && prefixes.stream().noneMatch(PyStringLiteralUtil::isRawPrefix)) {
-        for (Pair<TextRange, String> fragment : element.getDecodedFragments()) {
-          final String value = fragment.getSecond();
-          final int startOffset = fragment.getFirst().getStartOffset();
-          consumer.consumeToken(element, value, false, startOffset, TextRange.allOf(value), splitter);
+      for (PyStringElement stringElement : element.getStringElements()) {
+        final List<TextRange> literalPartRanges;
+        if (stringElement.isFormatted()) {
+          literalPartRanges = ((PyFormattedStringElement)stringElement).getLiteralPartRanges();
         }
-      }
-      else if (!prefixes.isEmpty()) {
-        for (TextRange valueTextRange : element.getStringValueTextRanges()) {
-          final String value = valueTextRange.substring(element.getText());
-          final int startOffset = valueTextRange.getStartOffset();
-          consumer.consumeToken(element, value, false, startOffset, TextRange.allOf(value), splitter);
+        else {
+          literalPartRanges = Collections.singletonList(stringElement.getContentRange());
+        }
+        final PyStringLiteralDecoder decoder = new PyStringLiteralDecoder(stringElement);
+        final boolean containsEscapes = stringElement.textContains('\\');
+        for (TextRange literalPartRange : literalPartRanges) {
+          final List<TextRange> escapeAwareRanges;
+          if (stringElement.isRaw() || !containsEscapes) {
+            escapeAwareRanges = Collections.singletonList(literalPartRange);
+          }
+          else {
+            escapeAwareRanges = ContainerUtil.map(decoder.decodeRange(literalPartRange), x -> x.getFirst());
+          }
+          for (TextRange escapeAwareRange : escapeAwareRanges) {
+            final String valueText = escapeAwareRange.substring(stringElement.getText());
+            consumer.consumeToken(stringElement, valueText, false, escapeAwareRange.getStartOffset(), TextRange.allOf(valueText), splitter);
+          }
         }
-      }
-      else {
-        consumer.consumeToken(element, splitter);
       }
     }
   }

python/testData/inspections/spelling/escapesInRawAndNormalGluedStringElements.py

@@ -0,0 +1 @@
+('\ncorrect' r'\<TYPO descr="Typo: In word 'ncorrect'">ncorrect</TYPO>')

python/testData/inspections/spelling/fStringExpression.py

@@ -0,0 +1,2 @@
+<TYPO descr="Typo: In word 'mstyped'">mstyped</TYPO> = 42
+f'{mstyped}'

python/testData/inspections/spelling/rawFString.py

@@ -0,0 +1 @@
+(f'\t<TYPO descr="Typo: In word 'eapot'">eapot</TYPO>{42}teapot' rf'\teapot{42}teapot')

python/testSrc/com/jetbrains/python/PySpellCheckerTest.java

@@ -48,13 +48,21 @@ public class PySpellCheckerTest extends PyTestCase {
     runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
   }
 
+  public void testFStringExpression() {
+    runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
+  }
+
+  public void testRawFString() {
+    runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
+  }
+
   // PY-20987
-  public void testGluedStringNodesAfterFirstWithPrefix() {
+  public void testEscapesInRawAndNormalGluedStringElements() {
     doTest();
   }
 
   // PY-20987
-  public void testGluedStringNodesAfterFirstRawWithBackslashes() {
+  public void testGluedStringNodesAfterFirstWithPrefix() {
     doTest();
   }
 

spellchecker/src/com/intellij/spellchecker/tokenizer/TokenConsumer.java

@@ -32,6 +32,14 @@ public abstract class TokenConsumer {
     consumeToken(element, text, useRename, 0, TextRange.allOf(text), splitter);
   }
 
+  /**
+   * @param element      PSI element on which problem descriptor will be set
+   * @param text         literal text that will be analyzed by spellchecker
+   * @param useRename    whether rename quick fix should be suggested instead of "change to"
+   * @param offset       offset inside element that serves as an anchor point for {@code rangeToCheck}
+   * @param rangeToCheck range text value corresponds to
+   * @param splitter
+   */
   public abstract void consumeToken(PsiElement element,
                                     String text,
                                     boolean useRename,