Initial implementation of introduce refactoring for substrings (PY-3654)

python/psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java

@@ -15,6 +15,7 @@ import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.ArrayUtil;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyTokenTypes;
+import com.jetbrains.python.PythonStringUtil;
 import com.jetbrains.python.psi.*;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -71,15 +72,42 @@ public class PyPsiUtils {
                                              @NotNull final PsiElement newExpression) {
     final Pair<PsiElement, TextRange> data = oldExpression.getUserData(SELECTION_BREAKS_AST_NODE);
     if (data != null) {
-      final PsiElement parent = data.first;
+      final PsiElement element = data.first;
       final TextRange textRange = data.second;
-      final String parentText = parent.getText();
+      final String parentText = element.getText();
       final String prefix = parentText.substring(0, textRange.getStartOffset());
-      final String suffix = parentText.substring(textRange.getEndOffset(), parent.getTextLength());
+      final String suffix = parentText.substring(textRange.getEndOffset(), element.getTextLength());
       final PyElementGenerator generator = PyElementGenerator.getInstance(oldExpression.getProject());
       final LanguageLevel languageLevel = LanguageLevel.forElement(oldExpression);
-      final PsiElement expression = generator.createFromText(languageLevel, parent.getClass(), prefix + newExpression.getText() + suffix);
-      return parent.replace(expression);
+      if (element instanceof PyStringLiteralExpression) {
+        final Pair<String, String> quotes = PythonStringUtil.getQuotes(parentText);
+        final PsiElement parent = element.getParent();
+        final boolean parensNeeded = parent instanceof PyExpression && !(parent instanceof PyParenthesizedExpression);
+        if (quotes != null) {
+          final String leftQuote = quotes.getFirst();
+          final String rightQuote = quotes.getSecond();
+          final StringBuilder builder = new StringBuilder();
+          if (parensNeeded) {
+            builder.append("(");
+          }
+          if (!leftQuote.endsWith(prefix)) {
+            builder.append(prefix + rightQuote + " + ");
+          }
+          final int pos = builder.toString().length();
+          builder.append(newExpression.getText());
+          if (!rightQuote.startsWith(suffix)) {
+            builder.append(" + " + leftQuote + suffix);
+          }
+          if (parensNeeded) {
+            builder.append(")");
+          }
+          final PsiElement expression = generator.createExpressionFromText(languageLevel, builder.toString());
+          final PsiElement newElement = element.replace(expression);
+          return newElement.findElementAt(pos);
+        }
+      }
+      final PsiElement expression = generator.createFromText(languageLevel, element.getClass(), prefix + newExpression.getText() + suffix);
+      return element.replace(expression);
     }
     else {
       return oldExpression.replace(newExpression);

python/src/com/jetbrains/python/refactoring/PyRefactoringUtil.java

@@ -42,6 +42,22 @@ public class PyRefactoringUtil {
           occurrences.add(element);
           return;
         }
+        if (element instanceof PyStringLiteralExpression) {
+          final Pair<PsiElement, TextRange> selection = pattern.getUserData(PyPsiUtils.SELECTION_BREAKS_AST_NODE);
+          if (selection != null) {
+            final String substring = selection.getSecond().substring(pattern.getText());
+            final PyStringLiteralExpression expr = (PyStringLiteralExpression)element;
+            final String text = element.getText();
+            if (text != null && expr.getStringNodes().size() == 1) {
+              final int start = text.indexOf(substring);
+              if (start >= 0) {
+                element.putUserData(PyPsiUtils.SELECTION_BREAKS_AST_NODE, Pair.create(element, TextRange.from(start, substring.length())));
+                occurrences.add(element);
+                return;
+              }
+            }
+          }
+        }
         element.acceptChildren(this);
       }
     };

python/src/com/jetbrains/python/refactoring/introduce/IntroduceHandler.java

@@ -29,6 +29,8 @@ import com.intellij.refactoring.util.CommonRefactoringUtil;
 import com.intellij.util.Function;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PyNames;
+import com.jetbrains.python.PythonStringUtil;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
@@ -44,22 +46,23 @@ import java.util.*;
 
 /**
  * @author Alexey.Ivanov
+ * @author vlan
  */
 abstract public class IntroduceHandler implements RefactoringActionHandler {
   protected static PsiElement findAnchor(List<PsiElement> occurrences) {
     PsiElement anchor = occurrences.get(0);
     next:
     do {
-      PyStatement statement = PsiTreeUtil.getParentOfType(anchor, PyStatement.class);
-
-      final PsiElement parent = statement.getParent();
-      for (PsiElement element : occurrences) {
-        if (!PsiTreeUtil.isAncestor(parent, element, true)) {
-          anchor = statement;
-          continue next;
+      final PyStatement statement = PsiTreeUtil.getParentOfType(anchor, PyStatement.class);
+      if (statement != null) {
+        final PsiElement parent = statement.getParent();
+        for (PsiElement element : occurrences) {
+          if (!PsiTreeUtil.isAncestor(parent, element, true)) {
+            anchor = statement;
+            continue next;
+          }
         }
       }
-
       return statement;
     }
     while (true);
@@ -112,7 +115,7 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
   protected PsiElement replaceExpression(PsiElement expression, PyExpression newExpression, IntroduceOperation operation) {
     PyExpressionStatement statement = PsiTreeUtil.getParentOfType(expression, PyExpressionStatement.class);
     if (statement != null) {
-      if (statement.getExpression() == expression) {
+      if (statement.getExpression() == expression && expression.getUserData(PyPsiUtils.SELECTION_BREAKS_AST_NODE) == null) {
         statement.delete();
         return null;
       }
@@ -169,6 +172,10 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
       }
     };
     String text = expression.getText();
+    final Pair<PsiElement, TextRange> selection = expression.getUserData(PyPsiUtils.SELECTION_BREAKS_AST_NODE);
+    if (selection != null) {
+      text = selection.getSecond().substring(text);
+    }
     if (expression instanceof PyCallExpression) {
       final PyExpression callee = ((PyCallExpression)expression).getCallee();
       if (callee != null) {
@@ -263,13 +270,26 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
       return;
     }
 
-    // Introduce refactoring for substrings is not supported yet
-    TextRange r = element1.getTextRange();
-    if (singleElementSelection && element1 instanceof PyStringLiteralExpression &&
-        (r.getStartOffset() < selectionModel.getSelectionStart() || r.getEndOffset() > selectionModel.getSelectionEnd())) {
-      showCannotPerformError(project, editor);
-      return;
+    if (singleElementSelection && element1 instanceof PyStringLiteralExpression) {
+      // TODO: Protect against escapes
+      // TODO: Protect against substrings with format characters
+      // TODO: Handle extracting substring from a string with formatting
+      final PyStringLiteralExpression literal = (PyStringLiteralExpression)element1;
+      if (literal.getStringNodes().size() > 1) {
+        showCannotPerformError(project, editor);
+        return;
+      }
+      final int offset = element1.getTextOffset();
+      final TextRange selectionRange = TextRange.create(selectionModel.getSelectionStart(), selectionModel.getSelectionEnd());
+      final TextRange elementRange = element1.getTextRange();
+      if (!elementRange.equals(selectionRange) && elementRange.contains(selectionRange)) {
+        final TextRange innerRange = literal.getStringValueTextRange();
+        final TextRange intersection = selectionRange.shiftRight(-offset).intersection(innerRange);
+        final TextRange finalRange = intersection != null ? intersection : selectionRange;
+        element1.putUserData(PyPsiUtils.SELECTION_BREAKS_AST_NODE, Pair.create(element1, finalRange));
+      }
     }
+
     if (!checkIntroduceContext(file, editor, element1)) {
       return;
     }
@@ -357,10 +377,10 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
                                     (PyExpression)element;
     operation.setInitializer(initializer);
 
-    if (initializer.getUserData(PyPsiUtils.SELECTION_BREAKS_AST_NODE) == null) {
+    if (initializer != null) {
       operation.setOccurrences(getOccurrences(element, initializer));
+      operation.setSuggestedNames(getSuggestedNames(initializer));
     }
-    operation.setSuggestedNames(getSuggestedNames(initializer));
     if (operation.getOccurrences().size() == 0) {
       operation.setReplaceAll(false);
     }
@@ -453,21 +473,33 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
 
     @Override
     public void visitPyStringLiteralExpression(PyStringLiteralExpression node) {
-      ASTNode child = node.getNode().getFirstChildNode();
-      while (child != null) {
-        String text = child.getText();
-        if (child.getElementType() == TokenType.WHITE_SPACE) {
-          if (text.contains("\n")) {
-            if (!text.contains("\\")) {
-              myResult.append("\\");
+      final Pair<PsiElement, TextRange> data = node.getUserData(PyPsiUtils.SELECTION_BREAKS_AST_NODE);
+      if (data != null) {
+        final PsiElement parent = data.getFirst();
+        final String text = parent.getText();
+        final Pair<String, String> detectedQuotes = PythonStringUtil.getQuotes(text);
+        final Pair<String, String> quotes = detectedQuotes != null ? detectedQuotes : Pair.create("'", "'");
+        final TextRange range = data.getSecond();
+        final String substring = range.substring(text);
+        myResult.append(quotes.getFirst() + substring + quotes.getSecond());
+      }
+      else {
+        ASTNode child = node.getNode().getFirstChildNode();
+        while (child != null) {
+          String text = child.getText();
+          if (child.getElementType() == TokenType.WHITE_SPACE) {
+            if (text.contains("\n")) {
+              if (!text.contains("\\")) {
+                myResult.append("\\");
+              }
+              myResult.append(text);
             }
+          }
+          else {
             myResult.append(text);
           }
+          child = child.getTreeNext();
         }
-        else {
-          myResult.append(text);
-        }
-        child = child.getTreeNext();
       }
     }
 
@@ -498,14 +530,7 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
   }
 
   protected List<PsiElement> getOccurrences(PsiElement element, @NotNull final PyExpression expression) {
-    PsiElement context = PsiTreeUtil.getParentOfType(expression, PyFunction.class);
-    if (context == null) {
-      context = PsiTreeUtil.getParentOfType(expression, PyClass.class);
-    }
-    if (context == null) {
-      context = expression.getContainingFile();
-    }
-    return PyRefactoringUtil.getOccurrences(expression, context);
+    return PyRefactoringUtil.getOccurrences(expression, ScopeUtil.getScopeOwner(expression));
   }
 
   private PsiElement performReplace(@NotNull final PsiElement declaration,
@@ -551,7 +576,7 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
   }
 
   protected PyExpression createExpression(Project project, String name, PsiElement declaration) {
-    return PyElementGenerator.getInstance(project).createExpressionFromText(name);
+    return PyElementGenerator.getInstance(project).createExpressionFromText(LanguageLevel.forElement(declaration), name);
   }
 
   @Nullable

python/testData/refactoring/introduceVariable/bytesSubstring.after.py

@@ -0,0 +1,2 @@
+a = b'bar'
+b'foo' + a + b'baz'

python/testData/refactoring/introduceVariable/bytesSubstring.py

@@ -0,0 +1 @@
+b'foo<selection>bar</selection>baz'

python/testData/refactoring/introduceVariable/leftQuoteSubstring.after.py

@@ -0,0 +1,2 @@
+a = "hello"
+print(a + " world")

python/testData/refactoring/introduceVariable/leftQuoteSubstring.py

@@ -0,0 +1 @@
+print(<selection>"hello</selection> world")

python/testData/refactoring/introduceVariable/leftSubstring.after.py

@@ -0,0 +1,2 @@
+a = "hello"
+print(a + " world")

python/testData/refactoring/introduceVariable/leftSubstring.py

@@ -0,0 +1 @@
+print("<selection>hello</selection> world")

python/testData/refactoring/introduceVariable/middleSubstring.after.py

@@ -0,0 +1,2 @@
+a = "lo wor"
+print("hel" + a + "ld")

python/testData/refactoring/introduceVariable/middleSubstring.py

@@ -0,0 +1 @@
+print("hel<selection>lo wor</selection>ld")

python/testData/refactoring/introduceVariable/rightSubstring.after.py

@@ -0,0 +1,2 @@
+a = "world"
+print("hello " + a)

python/testData/refactoring/introduceVariable/rightSubstring.py

@@ -0,0 +1 @@
+print("hello <selection>world</selection>")

python/testData/refactoring/introduceVariable/substringInExpression.after.py

@@ -0,0 +1,2 @@
+a = 'foo'
+print((a + 'bar').upper())

python/testData/refactoring/introduceVariable/substringInExpression.py

@@ -0,0 +1 @@
+print('<selection>foo</selection>bar'.upper())

python/testData/refactoring/introduceVariable/substringInExpressionStatement.after.py

@@ -0,0 +1,2 @@
+a = "two"
+"one " + a + " three"

python/testData/refactoring/introduceVariable/substringInExpressionStatement.py

@@ -0,0 +1 @@
+"one <selection>two</selection> three"

python/testData/refactoring/introduceVariable/substringInStatement.after.py

@@ -0,0 +1,2 @@
+a = 'foo'
+x = a + 'bar'

python/testData/refactoring/introduceVariable/substringInStatement.py

@@ -0,0 +1 @@
+x = '<selection>foo</selection>bar'

python/testData/refactoring/introduceVariable/tripleQuotedSubstring.after.py

@@ -0,0 +1,6 @@
+a = """Three"""
+print(""""One two
+
+* """ + a + """
+* Four
+* Five""")

python/testData/refactoring/introduceVariable/tripleQuotedSubstring.py

@@ -0,0 +1,5 @@
+print(""""One two
+
+* <selection>Three</selection>
+* Four
+* Five""")

python/testSrc/com/jetbrains/python/refactoring/PyIntroduceVariableTest.java

@@ -88,6 +88,51 @@ public class PyIntroduceVariableTest extends PyIntroduceTestCase {
     doTest();
   }
 
+  // PY-3654
+  public void testLeftSubstring() {
+    doTest();
+  }
+
+  // PY-3654
+  public void testRightSubstring() {
+    doTest();
+  }
+
+  // PY-3654
+  public void testMiddleSubstring() {
+    doTest();
+  }
+
+  // PY-3654
+  public void testLeftQuoteSubstring() {
+    doTest();
+  }
+
+  // PY-3654
+  public void testSubstringInExpression() {
+    doTest();
+  }
+
+  // PY-3654
+  public void testSubstringInStatement() {
+    doTest();
+  }
+
+  // PY-3654
+  public void testTripleQuotedSubstring() {
+    doTest();
+  }
+
+  // PY-3654
+  public void testSubstringInExpressionStatement() {
+    doTest();
+  }
+
+  // PY-6354
+  public void testBytesSubstring() {
+    doTest();
+  }
+
   private void doTestCannotPerform() {
     boolean thrownExpectedException = false;
     try {