PY-8991, PY-8990 RemovePrefixQuickfix handles glued string literals

python/src/com/jetbrains/python/inspections/quickfix/RemovePrefixQuickFix.java

@@ -17,14 +17,16 @@ package com.jetbrains.python.inspections.quickfix;
 
 import com.intellij.codeInspection.LocalQuickFix;
 import com.intellij.codeInspection.ProblemDescriptor;
+import com.intellij.lang.ASTNode;
 import com.intellij.openapi.project.Project;
-import com.intellij.psi.PsiElement;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.psi.PyElementGenerator;
 import com.jetbrains.python.psi.PyStringLiteralExpression;
 import com.jetbrains.python.psi.impl.PyStringLiteralExpressionImpl;
 import org.jetbrains.annotations.NotNull;
 
+import static com.jetbrains.python.psi.PyUtil.as;
+
 /**
  * Created by IntelliJ IDEA.
  * Author: Alexey.Ivanov
@@ -39,24 +41,30 @@ public class RemovePrefixQuickFix implements LocalQuickFix {
   }
 
   @NotNull
-
   @Override
   public String getName() {
     return PyBundle.message("INTN.remove.leading.$0", myPrefix);
   }
 
   @NotNull
+  @Override
   public String getFamilyName() {
     return PyBundle.message("INTN.remove.leading.prefix");
   }
 
   @Override
   public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
-    PsiElement stringLiteralExpression = descriptor.getPsiElement();
-    if (stringLiteralExpression instanceof PyStringLiteralExpression) {
-      PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
-      final int length = PyStringLiteralExpressionImpl.getPrefixLength(stringLiteralExpression.getText());
-      stringLiteralExpression.replace(elementGenerator.createExpressionFromText(stringLiteralExpression.getText().substring(length)));
+    final PyStringLiteralExpression pyString = as(descriptor.getPsiElement(), PyStringLiteralExpression.class);
+    if (pyString != null) {
+      final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
+      for (ASTNode node : pyString.getStringNodes()) {
+        final String nodeText = node.getText();
+        final int prefixLength = PyStringLiteralExpressionImpl.getPrefixLength(nodeText);
+        if (nodeText.substring(0, prefixLength).equalsIgnoreCase(myPrefix)) {
+          final PyStringLiteralExpression replacement = elementGenerator.createStringLiteralAlreadyEscaped(nodeText.substring(prefixLength));
+          node.getPsi().replace(replacement.getFirstChild());
+        }
+      }
     }
   }
 }

python/src/com/jetbrains/python/validation/CompatibilityVisitor.java

@@ -23,11 +23,11 @@ import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.inspections.quickfix.*;
 import com.jetbrains.python.psi.*;
-import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.psi.impl.PyStringLiteralExpressionImpl;
 import org.jetbrains.annotations.Nullable;
 
@@ -241,17 +241,17 @@ public abstract class CompatibilityVisitor extends PyAnnotator {
   @Override
   public void visitPyStringLiteralExpression(final PyStringLiteralExpression node) {
     super.visitPyStringLiteralExpression(node);
-    List<ASTNode> stringNodes = node.getStringNodes();
+    final List<ASTNode> stringNodes = node.getStringNodes();
 
     for (ASTNode stringNode : stringNodes) {
       int len = 0;
-      StringBuilder message = new StringBuilder(myCommonMessage);
-      String nodeText = stringNode.getText();
-      int index = PyStringLiteralExpressionImpl.getPrefixLength(nodeText);
-      String prefix = nodeText.substring(0, index).toUpperCase();
+      final StringBuilder message = new StringBuilder(myCommonMessage);
+      final String nodeText = stringNode.getText();
+      final int index = PyStringLiteralExpressionImpl.getPrefixLength(nodeText);
+      final String prefix = nodeText.substring(0, index).toUpperCase();
       final TextRange range = TextRange.create(stringNode.getStartOffset(), stringNode.getStartOffset() + index);
       for (int i = 0; i != myVersionsToProcess.size(); ++i) {
-        LanguageLevel languageLevel = myVersionsToProcess.get(i);
+        final LanguageLevel languageLevel = myVersionsToProcess.get(i);
         if (prefix.isEmpty()) continue;
 
         final Set<String> prefixes = AVAILABLE_PREFIXES.get(languageLevel);

python/testData/inspections/RemoveUnicodePrefixFromGluedStringNodesInParenthesis.py

@@ -0,0 +1,3 @@
+s = (<error descr="Python version 3.2 does not support a 'U' prefix">u<caret></error>"string \n"
+    <error descr="Python version 3.2 does not support a 'U' prefix">u</error>"next line"
+)

python/testData/inspections/RemoveUnicodePrefixFromGluedStringNodesInParenthesis_after.py

@@ -0,0 +1,3 @@
+s = ("string \n"
+     "next line"
+)

python/testData/inspections/RemoveUnicodePrefixFromGluedStringNodesWithSlash.py

@@ -0,0 +1,2 @@
+s = <error descr="Python version 3.2 does not support a 'U' prefix">u<caret></error>"string\n" \
+    <error descr="Python version 3.2 does not support a 'U' prefix">u</error>"next line"

python/testData/inspections/RemoveUnicodePrefixFromGluedStringNodesWithSlash_after.py

@@ -0,0 +1,2 @@
+s = "string\n" \
+    "next line"

python/testSrc/com/jetbrains/python/PyQuickFixTest.java

@@ -470,6 +470,34 @@ public class PyQuickFixTest extends PyTestCase {
     myFixture.checkResultByFile(graftBeforeExt(fileName, "_after"));
   }
 
+  // PY-8991
+  public void testRemoveUnicodePrefixFromGluedStringNodesWithSlash() {
+    runWithLanguageLevel(LanguageLevel.PYTHON32, new Runnable() {
+      public void run() {
+        myFixture.configureByFiles(getTestDataPath() + getTestName(false) + ".py");
+        myFixture.checkHighlighting(true, false, false);
+        final IntentionAction intentionAction = myFixture.findSingleIntention(PyBundle.message("INTN.remove.leading.$0", "U"));
+        assertNotNull(intentionAction);
+        myFixture.launchAction(intentionAction);
+        myFixture.checkResultByFile(getTestName(false) + "_after.py");
+      }
+    });
+  }
+
+  // PY-8990
+  public void testRemoveUnicodePrefixFromGluedStringNodesInParenthesis() {
+    runWithLanguageLevel(LanguageLevel.PYTHON32, new Runnable() {
+      public void run() {
+        myFixture.configureByFiles(getTestDataPath() + getTestName(false) + ".py");
+        myFixture.checkHighlighting(true, false, false);
+        final IntentionAction intentionAction = myFixture.findSingleIntention(PyBundle.message("INTN.remove.leading.$0", "U"));
+        assertNotNull(intentionAction);
+        myFixture.launchAction(intentionAction);
+        myFixture.checkResultByFile(getTestName(false) + "_after.py");
+      }
+    });
+  }
+
   @Override
   @NonNls
   protected String getTestDataPath() {