PY-20777 Disable commenting inside injected Python expressions in f-strings

This solution uses a custom key to mark an injected file as unsuitable for line
commenting, but we should rather generalize it in some other way instead of
further hardcoding language-specific string prefixes in CommentByLineCommentHandler.

Reviewed in IDEA-CR-35030.

platform/lang-impl/src/com/intellij/codeInsight/generation/CommentByLineCommentHandler.java

@@ -22,6 +22,7 @@ import com.intellij.openapi.fileTypes.impl.AbstractFileType;
 import com.intellij.openapi.fileTypes.impl.CustomSyntaxTableFileType;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Comparing;
+import com.intellij.openapi.util.Key;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.registry.Registry;
 import com.intellij.openapi.util.text.StringUtil;
@@ -45,6 +46,21 @@ import java.util.Map;
 
 public class CommentByLineCommentHandler extends MultiCaretCodeInsightActionHandler {
 
+  private static final Key<Boolean> INJECTION_FORBIDS_LINE_COMMENTS = Key.create("INJECTION_FORBIDS_LINE_COMMENTS");
+
+  /**
+   * Disable line commenting in an injected file making this action operate on its host file instead.
+   *
+   * @param file injected file where line comment action shouldn't be available
+   */
+  public static void markInjectedFileUnsuitableForLineComment(@NotNull PsiFile file) {
+    if (!InjectedLanguageManager.getInstance(file.getProject()).isInjectedFragment(file)) {
+      throw new IllegalArgumentException("This method should be called only on injected files");
+    }
+
+    file.putUserData(INJECTION_FORBIDS_LINE_COMMENTS, true);
+  }
+
   private final List<Block> myBlocks = new ArrayList<>();
 
   @Override
@@ -54,13 +70,10 @@ public class CommentByLineCommentHandler extends MultiCaretCodeInsightActionHand
 
     PsiElement context = InjectedLanguageManager.getInstance(file.getProject()).getInjectionHost(file);
 
-    if (context != null && (context.textContains('\'') || context.textContains('\"') || context.textContains('/'))) {
-      String s = context.getText();
-      if (StringUtil.startsWith(s, "\"") || StringUtil.startsWith(s, "\'") || StringUtil.startsWith(s, "/")) {
-        file = context.getContainingFile();
-        editor = editor instanceof EditorWindow ? ((EditorWindow)editor).getDelegate() : editor;
-        caret = caret instanceof InjectedCaret ? ((InjectedCaret)caret).getDelegate() : caret;
-      }
+    if (context != null && shouldCommentInHostFile(file, context)) {
+      file = context.getContainingFile();
+      editor = editor instanceof EditorWindow ? ((EditorWindow)editor).getDelegate() : editor;
+      caret = caret instanceof InjectedCaret ? ((InjectedCaret)caret).getDelegate() : caret;
     }
 
     Document document = editor.getDocument();
@@ -127,6 +140,17 @@ public class CommentByLineCommentHandler extends MultiCaretCodeInsightActionHand
                                wholeLinesSelected ? CaretUpdate.RESTORE_SELECTION : null;
     }
 
+  private static boolean shouldCommentInHostFile(@NotNull PsiFile file, @NotNull PsiElement context) {
+    if (file.getUserData(INJECTION_FORBIDS_LINE_COMMENTS) != null) {
+      return true;
+    }
+    if (context.textContains('\'') || context.textContains('\"') || context.textContains('/')) {
+      final String s = context.getText();
+      return StringUtil.startsWith(s, "\"") || StringUtil.startsWith(s, "\'") || StringUtil.startsWith(s, "/");
+    }
+    return false;
+  }
+
   @Override
   public void postInvoke() {
     FeatureUsageTracker.getInstance().triggerFeatureUsed("codeassists.comment.line");

python/src/com/jetbrains/python/codeInsight/fstrings/PyFStringsInjector.java

@@ -15,15 +15,20 @@
  */
 package com.jetbrains.python.codeInsight.fstrings;
 
+import com.intellij.codeInsight.generation.CommentByLineCommentHandler;
 import com.intellij.lang.ASTNode;
 import com.intellij.lang.Language;
+import com.intellij.lang.injection.InjectedLanguageManager;
 import com.intellij.lang.injection.MultiHostRegistrar;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.jetbrains.python.codeInsight.PyInjectorBase;
 import com.jetbrains.python.codeInsight.fstrings.FStringParser.Fragment;
 import com.jetbrains.python.documentation.doctest.PyDocstringLanguageDialect;
 import com.jetbrains.python.psi.PyStringLiteralExpression;
 import com.jetbrains.python.psi.PyUtil.StringNodeInfo;
+import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -47,15 +52,31 @@ public class PyFStringsInjector extends PyInjectorBase {
   }
 
   public static void injectFStringFragments(@NotNull MultiHostRegistrar registrar, @NotNull PyStringLiteralExpression pyString) {
+    final PyDocstringLanguageDialect docstringLanguage = PyDocstringLanguageDialect.getInstance();
     for (ASTNode node : pyString.getStringNodes()) {
       final int relNodeOffset = node.getTextRange().getStartOffset() - pyString.getTextRange().getStartOffset();
       for (Fragment offsets : getInjectionRanges(node)) {
         if (offsets.containsNamedUnicodeEscape()) continue;
-        registrar.startInjecting(PyDocstringLanguageDialect.getInstance());
+        registrar.startInjecting(docstringLanguage);
         registrar.addPlace(null, null, pyString, offsets.getContentRange().shiftRight(relNodeOffset));
         registrar.doneInjecting();
       }
     }
+
+    disableCommentingInFragments(pyString);
+  }
+
+  private static void disableCommentingInFragments(@NotNull PyStringLiteralExpression pyString) {
+    final Project project = pyString.getProject();
+    final InjectedLanguageManager injectedLanguageManager = InjectedLanguageManager.getInstance(project);
+    final PsiDocumentManager documentManager = PsiDocumentManager.getInstance(project);
+    final PyDocstringLanguageDialect docstringLanguage = PyDocstringLanguageDialect.getInstance();
+
+    StreamEx.of(injectedLanguageManager.getCachedInjectedDocumentsInRange(pyString.getContainingFile(), pyString.getTextRange()))
+      .map(window -> documentManager.getPsiFile(window))
+      .nonNull()
+      .filter(file -> file.getLanguage().isKindOf(docstringLanguage))
+      .forEach(CommentByLineCommentHandler::markInjectedFileUnsuitableForLineComment);
   }
 
   @NotNull

python/testData/commenter/lineCommentInFStringFragment.py

@@ -0,0 +1 @@
+s = f'{foo} {b<caret>ar}'

python/testData/commenter/lineCommentInFStringFragment_after.py

@@ -0,0 +1 @@
+# s = f'{foo} {bar}'

python/testSrc/com/jetbrains/python/PyCommenterTest.java

@@ -17,6 +17,7 @@ package com.jetbrains.python;
 
 import com.intellij.openapi.actionSystem.IdeActions;
 import com.jetbrains.python.fixtures.PyTestCase;
+import com.jetbrains.python.psi.LanguageLevel;
 
 /**
  * @author yole
@@ -30,6 +31,11 @@ public class PyCommenterTest extends PyTestCase {
     doTest();
   }
 
+  // PY-20777
+  public void testLineCommentInFStringFragment() {
+    runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
+  }
+
   private void doTest() {
     myFixture.configureByFile("commenter/" + getTestName(true) + ".py");
     myFixture.performEditorAction(IdeActions.ACTION_COMMENT_LINE);