PY-42700 Offer f-string-like completion after "{" in non-interpolated strings

The main idea here is that we also automatically convert the containing string
literal into a proper f-string upon selecting one of such variants.

Hopefully, it addresses the common problem of forgetting to add "f" prefix
before an interpolated (in mind) string literal in a more discoverable fashion
than a dedicated intention or a postfix template would.

GitOrigin-RevId: f59065abeb4a3b0cd7a0de42c9c06c9fad318b19

python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml

@@ -102,6 +102,8 @@
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyClassNameCompletionContributor"/>
     <completion.contributor language="Python"
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyModulePackageCompletionContributor"/>
+    <completion.contributor language="Python"
+                            implementationClass="com.jetbrains.python.codeInsight.completion.PyFStringLikeCompletionContributor"/>
 
     <lang.tokenSeparatorGenerator language="Python" implementationClass="com.jetbrains.python.PyTokenSeparatorGenerator"/>
 

python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyFStringLikeCompletionContributor.java

@@ -0,0 +1,105 @@
+package com.jetbrains.python.codeInsight.completion;
+
+import com.intellij.codeInsight.completion.*;
+import com.intellij.codeInsight.lookup.LookupElement;
+import com.intellij.codeInsight.lookup.LookupElementDecorator;
+import com.intellij.featureStatistics.FeatureUsageTracker;
+import com.intellij.openapi.editor.Document;
+import com.intellij.patterns.PsiElementPattern;
+import com.intellij.psi.PsiReference;
+import com.intellij.util.ProcessingContext;
+import com.intellij.util.containers.ContainerUtil;
+import com.intellij.util.text.CharArrayUtil;
+import com.jetbrains.python.PyNames;
+import com.jetbrains.python.psi.*;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.List;
+
+import static com.intellij.patterns.PlatformPatterns.psiElement;
+import static com.jetbrains.python.psi.PyUtil.as;
+
+/**
+ * Provides completion results after "{" in ordinary non-formatted string literals as if that literal
+ * was a proper f-string and adds a missing "f" prefix and a closing brace on selection of such a variant.
+ * <p>
+ * For instance, in the next fragment
+ * <pre><code>
+ * for line in f:
+ *     print('# {li&lt;caret&gt;')
+ * </code></pre>
+ * selecting "line" from the list of completion suggestions will automatically transform it into
+ * <pre><code>
+ * for line in f:
+ *     print(f'# {line&lt;caret&gt;}')
+ * </code></pre>
+ */
+public class PyFStringLikeCompletionContributor extends CompletionContributor {
+  private static final String FEATURE_ID = "python.completion.fstring.like";
+
+  private static final PsiElementPattern.Capture<PyPlainStringElement> INSIDE_NON_FORMATTED_STRING_ELEMENT =
+    psiElement(PyPlainStringElement.class)
+    .andNot(psiElement().inside(PyStringFormatCompletionContributor.FORMAT_STRING_CAPTURE));
+
+  public PyFStringLikeCompletionContributor() {
+    extend(CompletionType.BASIC, INSIDE_NON_FORMATTED_STRING_ELEMENT, new CompletionProvider<CompletionParameters>() {
+      @Override
+      protected void addCompletions(@NotNull CompletionParameters parameters,
+                                    @NotNull ProcessingContext context,
+                                    @NotNull CompletionResultSet result) {
+        PyPlainStringElement stringElem = as(parameters.getPosition(), PyPlainStringElement.class);
+        if (stringElem == null || stringElem.isBytes() || stringElem.isUnicode()) {
+          return;
+        }
+        if (LanguageLevel.forElement(stringElem).isOlderThan(LanguageLevel.PYTHON36)) {
+          return;
+        }
+        PyStringLiteralExpression stringLiteral = (PyStringLiteralExpression)stringElem.getParent();
+        String stringElemText = stringElem.getText();
+        int offset = parameters.getOffset();
+        int stringElemStart = stringElem.getTextRange().getStartOffset();
+        int relOffset = offset - stringElemStart;
+        int braceOffset = CharArrayUtil.shiftBackwardUntil(stringElemText, relOffset - 1, "{");
+        if (braceOffset < 0) {
+          return;
+        }
+        String completionPrefix = stringElemText.substring(braceOffset + 1, relOffset);
+        if (!PyNames.isIdentifier(completionPrefix)) {
+          return;
+        }
+        PyExpression fString = PyUtil.createExpressionFromFragment("f" + stringElemText, stringLiteral.getParent());
+        assert fString != null;
+        PsiReference reference = fString.findReferenceAt(relOffset + 1);
+        if (reference == null) {
+          return;
+        }
+        List<@NotNull LookupElement> fStringVariants = ContainerUtil.mapNotNull(reference.getVariants(), v -> as(v, LookupElement.class));
+        if (fStringVariants.isEmpty()) {
+          return;
+        }
+        CompletionResultSet prefixPatchedResultSet = result.withPrefixMatcher(completionPrefix);
+        for (LookupElement variant : fStringVariants) {
+          prefixPatchedResultSet.addElement(new LookupElementDecorator<LookupElement>(variant) {
+            @Override
+            public void handleInsert(@NotNull InsertionContext context) {
+              FeatureUsageTracker.getInstance().triggerFeatureUsed(FEATURE_ID);
+
+              super.handleInsert(context);
+              Document document = context.getDocument();
+              CharSequence docChars = document.getCharsSequence();
+              int tailOffset = context.getTailOffset();
+              if (!(tailOffset < document.getTextLength() && docChars.charAt(tailOffset) == '}')) {
+                document.insertString(tailOffset, "}");
+              }
+              // It can happen when completion is invoked on multiple carets inside the same string
+              String stringElemPrefix = PyStringLiteralUtil.getPrefix(docChars, stringElemStart);
+              if (!PyStringLiteralUtil.isFormattedPrefix(stringElemPrefix)) {
+                document.insertString(stringElemStart, "f");
+              }
+            }
+          });
+        }
+      }
+    });
+  }
+}

python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyStringFormatCompletionContributor.java

@@ -48,7 +48,7 @@ public class PyStringFormatCompletionContributor extends CompletionContributor {
       }
     };
 
-  private static final PsiElementPattern.Capture<PyStringLiteralExpression> FORMAT_STRING_CAPTURE =
+  public static final PsiElementPattern.Capture<PyStringLiteralExpression> FORMAT_STRING_CAPTURE =
     psiElement(PyStringLiteralExpression.class)
       .withParent(psiElement(PyReferenceExpression.class).with(FORMAT_CALL_PATTERN_CONDITION))
       .withSuperParent(2, PyCallExpression.class);

python/testData/completion/fStringLikeCompletionAvailableAfterOpeningBraceFollowingEscapedOne.after.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = f'foo{{{my_expr}'

python/testData/completion/fStringLikeCompletionAvailableAfterOpeningBraceFollowingEscapedOne.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = 'foo{{{my_e<caret>'

python/testData/completion/fStringLikeCompletionDoesNotDuplicateClosingBrace.after.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = f'foo{my_expr}'

python/testData/completion/fStringLikeCompletionDoesNotDuplicateClosingBrace.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = 'foo{my_e<caret>}'

python/testData/completion/fStringLikeCompletionDoesNotWorkInStringWithInjections.py

@@ -0,0 +1,3 @@
+my_expr = True
+# language=regexp
+pattern = r'{my_e<caret>'

python/testData/completion/fStringLikeCompletionInOrdinaryStringLiterals.after.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = f'foo{my_expr}'

python/testData/completion/fStringLikeCompletionInOrdinaryStringLiterals.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = 'foo{my_e<caret>'

python/testData/completion/fStringLikeCompletionNotAvailableAfterEscapedOpeningBrace.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = 'foo{{my_e<caret>'

python/testData/completion/fStringLikeCompletionNotAvailableBefore36.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = 'foo{my_e<caret>'

python/testData/completion/fStringLikeCompletionNotAvailableInByteLiterals.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = b'foo{my_e<caret>'

python/testData/completion/fStringLikeCompletionNotAvailableInStrFormatCalls.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = 'foo{my_e<caret>'.format(my_expr='spam')

python/testData/completion/fStringLikeCompletionNotAvailableInUnicodeLiterals.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = u'foo{my_e<caret>'

python/testData/completion/fStringLikeCompletionOnMultipleCaretsDoesNotDuplicatePrefix.after.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = f'foo{my_expr} bar{my_expr}'

python/testData/completion/fStringLikeCompletionOnMultipleCaretsDoesNotDuplicatePrefix.py

@@ -0,0 +1,2 @@
+my_expr = 42
+s = 'foo{my_e<caret> bar{my_e<caret>'

python/testData/completion/fStringLikeCompletionPreservesParenthesesForCallables.after.py

@@ -0,0 +1 @@
+s = f'foo{oct(<caret>)}'

python/testData/completion/fStringLikeCompletionPreservesParenthesesForCallables.py

@@ -0,0 +1 @@
+s = 'foo{oct<caret>'

python/testSrc/com/jetbrains/python/Py3CompletionTest.java

@@ -2,6 +2,7 @@
 package com.jetbrains.python;
 
 import com.intellij.codeInsight.completion.impl.CamelHumpMatcher;
+import com.intellij.codeInsight.lookup.LookupElement;
 import com.intellij.openapi.module.Module;
 import com.intellij.testFramework.LightProjectDescriptor;
 import com.intellij.testFramework.PsiTestUtil;
@@ -65,6 +66,14 @@ public class Py3CompletionTest extends PyTestCase {
     myFixture.checkResultByFile(testName + ".after.py");
   }
 
+  public void doNegativeTest() {
+    final String testName = getTestName(true);
+    myFixture.configureByFile(testName + ".py");
+    LookupElement[] variants = myFixture.completeBasic();
+    assertNotNull("Expected no completion variants, but one item was auto-completed", variants);
+    assertEmpty(variants);
+  }
+
   private void doMultiFileTest() {
     myFixture.copyDirectoryToProject(getTestName(true), "");
     myFixture.configureByFile("a.py");
@@ -408,6 +417,61 @@ public class Py3CompletionTest extends PyTestCase {
     myFixture.checkResultByFile(getTestName(true) + "/a.after.py");
   }
 
+  // PY-42700
+  public void testFStringLikeCompletionInOrdinaryStringLiterals() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doTest);
+  }
+
+  // PY-42700
+  public void testFStringLikeCompletionPreservesParenthesesForCallables() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doTest);
+  }
+
+  // PY-42700
+  public void testFStringLikeCompletionNotAvailableBefore36() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, this::doNegativeTest);
+  }
+
+  // PY-42700
+  public void testFStringLikeCompletionNotAvailableInByteLiterals() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doNegativeTest);
+  }
+
+  // PY-42700
+  public void testFStringLikeCompletionNotAvailableInUnicodeLiterals() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doNegativeTest);
+  }
+
+  // PY-42700
+  public void testFStringLikeCompletionNotAvailableInStrFormatCalls() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doNegativeTest);
+  }
+
+  // PY-42700
+  public void testFStringLikeCompletionNotAvailableAfterEscapedOpeningBrace() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doNegativeTest);
+  }
+
+  // PY-42700
+  public void testFStringLikeCompletionAvailableAfterOpeningBraceFollowingEscapedOne() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doTest);
+  }
+
+  // PY-42700
+  public void testFStringLikeCompletionDoesNotDuplicateClosingBrace() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doTest);
+  }
+
+  // PY-42700
+  public void testFStringLikeCompletionOnMultipleCaretsDoesNotDuplicatePrefix() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doTest);
+  }
+
+  // PY-42700
+  public void testFStringLikeCompletionDoesNotWorkInStringWithInjections() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doNegativeTest);
+  }
+
   @Override
   protected String getTestDataPath() {
     return super.getTestDataPath() + "/completion";