From 4100824a3e1ad45c4be6bb887ed06cc1ee29d2ca Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Mon, 10 Sep 2018 19:12:33 +0300
Subject: [PATCH] PY-31442 Restore functioning of Join Lines for f-strings

It's not able to handle more complex cases yet, though. Say, if we want to join plain
and formatted strings literals.
---
 .../python/editor/PyJoinLinesHandler.java     | 75 ++++++++++---------
 .../joinLines/FStringAndPlainLiteral-after.py |  1 +
 .../joinLines/FStringAndPlainLiteral.py       |  2 +
 python/testData/joinLines/FStrings-after.py   |  1 +
 python/testData/joinLines/FStrings.py         |  2 +
 .../com/jetbrains/python/PyJoinLinesTest.java |  8 ++
 6 files changed, 52 insertions(+), 37 deletions(-)
 create mode 100644 python/testData/joinLines/FStringAndPlainLiteral-after.py
 create mode 100644 python/testData/joinLines/FStringAndPlainLiteral.py
 create mode 100644 python/testData/joinLines/FStrings-after.py
 create mode 100644 python/testData/joinLines/FStrings.py

diff --git a/python/src/com/jetbrains/python/editor/PyJoinLinesHandler.java b/python/src/com/jetbrains/python/editor/PyJoinLinesHandler.java
index e3629be66f51..1e84eec16de0 100644
--- a/python/src/com/jetbrains/python/editor/PyJoinLinesHandler.java
+++ b/python/src/com/jetbrains/python/editor/PyJoinLinesHandler.java
@@ -33,8 +33,6 @@ import org.jetbrains.annotations.Nullable;
 
 import java.util.List;
 
-import static com.jetbrains.python.psi.PyUtil.StringNodeInfo;
-
 /**
  * Joins lines sanely.
  * <ul>
@@ -226,47 +224,49 @@ public class PyJoinLinesHandler implements JoinRawLinesHandlerDelegate {
   private static class StringLiteralJoiner implements Joiner {
     @Override
     public Result join(@NotNull Request req) {
-      if (req.leftElem != req.rightElem) {
-        final PsiElement parent = req.rightElem.getParent();
-        if ((req.leftElem.getParent() == parent && parent instanceof PyStringLiteralExpression) ||
+      final PyStringElement leftStringElem = PsiTreeUtil.getParentOfType(req.leftElem, PyStringElement.class, false);
+      final PyStringElement rightStringElem = PsiTreeUtil.getParentOfType(req.rightElem, PyStringElement.class, false);
+      if (leftStringElem != null && rightStringElem != null && leftStringElem != rightStringElem) {
+        final PsiElement parent = rightStringElem.getParent();
+        if ((leftStringElem.getParent() == parent && parent instanceof PyStringLiteralExpression) ||
             (req.leftExpr instanceof PyStringLiteralExpression && req.rightExpr instanceof PyStringLiteralExpression)) {
-          final StringNodeInfo leftNodeInfo = new StringNodeInfo(req.leftElem);
-          final StringNodeInfo rightNodeInfo = new StringNodeInfo(req.rightElem);
-
-          if (leftNodeInfo.isTerminated() && rightNodeInfo.isTerminated()) {
-            int quotesMaxLength = Math.max(leftNodeInfo.getQuote().length(), rightNodeInfo.getQuote().length());
+          
+          if (leftStringElem.isTerminated() && rightStringElem.isTerminated() && haveSamePrefixes(leftStringElem, rightStringElem)) {
+            final String leftElemQuotes = leftStringElem.getQuote();
+            final String rightElemQuotes = rightStringElem.getQuote();
+            int quotesMaxLength = Math.max(leftElemQuotes.length(), rightElemQuotes.length());
             int stringToJoinMaxLength = getStringToJoinMaxLength(req, quotesMaxLength);
-            final String replacement = findReplacement(rightNodeInfo.getContent(), stringToJoinMaxLength);
+            final String replacement = findReplacement(rightStringElem.getContent(), stringToJoinMaxLength);
 
-            if (leftNodeInfo.equals(rightNodeInfo)) {
-              return getResultAndSplitStringIfTooLong(req, leftNodeInfo, rightNodeInfo, replacement, leftNodeInfo.getQuote());
+            if (leftElemQuotes.equals(rightElemQuotes)) {
+              return getResultAndSplitStringIfTooLong(req, leftStringElem, rightStringElem, replacement, leftElemQuotes);
             }
 
-            return processStringsWithDifferentQuotes(req, leftNodeInfo, rightNodeInfo, replacement);
+            return processStringsWithDifferentQuotes(req, leftStringElem, rightStringElem, replacement);
           }
         }
       }
       return null;
     }
 
+
     @Nullable
     private static Result processStringsWithDifferentQuotes(@NotNull final Request req,
-                                                            @NotNull final StringNodeInfo leftNodeInfo,
-                                                            @NotNull final StringNodeInfo rightNodeInfo,
+                                                            @NotNull final PyStringElement leftElem,
+                                                            @NotNull final PyStringElement rightElem,
                                                             @NotNull final String replacement) {
-      if (haveSamePrefixes(leftNodeInfo, rightNodeInfo) && !leftNodeInfo.isTripleQuoted() && !rightNodeInfo.isTripleQuoted()) {
-
-        if (!rightNodeInfo.getContent().contains(leftNodeInfo.getQuote())) {
-          final int quotePos = rightNodeInfo.getAbsoluteContentRange().getEndOffset();
-          final String quote = leftNodeInfo.getQuote();
+      if (!leftElem.isTripleQuoted() && !rightElem.isTripleQuoted()) {
+        if (!rightElem.getContent().contains(leftElem.getQuote())) {
+          final int quotePos = rightElem.getTextOffset() + rightElem.getContentRange().getEndOffset();
+          final String quote = leftElem.getQuote();
           req.document.replaceString(quotePos, quotePos + 1, quote);
-          return getResultAndSplitStringIfTooLong(req, leftNodeInfo, rightNodeInfo, replacement, quote);
+          return getResultAndSplitStringIfTooLong(req, leftElem, rightElem, replacement, quote);
         }
-        else if (!leftNodeInfo.getContent().contains(rightNodeInfo.getQuote())) {
-          final int quotePos = leftNodeInfo.getAbsoluteContentRange().getStartOffset() - 1;
-          final String quote = rightNodeInfo.getQuote();
+        else if (!leftElem.getContent().contains(rightElem.getQuote())) {
+          final int quotePos = leftElem.getTextOffset() + leftElem.getContentRange().getStartOffset() - 1;
+          final String quote = rightElem.getQuote();
           req.document.replaceString(quotePos, quotePos + 1, quote);
-          return getResultAndSplitStringIfTooLong(req, leftNodeInfo, rightNodeInfo, replacement, quote);
+          return getResultAndSplitStringIfTooLong(req, leftElem, rightElem, replacement, quote);
         }
       }
       return null;
@@ -274,28 +274,29 @@ public class PyJoinLinesHandler implements JoinRawLinesHandlerDelegate {
 
     @NotNull
     private static Result getResultAndSplitStringIfTooLong(@NotNull final Request req,
-                                                           @NotNull final StringNodeInfo leftNodeInfo,
-                                                           @NotNull final StringNodeInfo rightNodeInfo,
+                                                           @NotNull final PyStringElement leftElem,
+                                                           @NotNull final PyStringElement rightElem,
                                                            @NotNull final String replacement,
                                                            @NotNull final String quote) {
-      int cutIntoRight = rightNodeInfo.getContentRange().getStartOffset();
+      int cutIntoRight = rightElem.getContentRange().getStartOffset();
       String lineEnd = "";
       if (!replacement.isEmpty()) {
-        cutIntoRight = replacement.length() + rightNodeInfo.getQuote().length();
-        int contentWithQuoteStartColumn = leftNodeInfo.getAbsoluteContentRange().getStartOffset() - getLeftLineStartOffset(req)
-                                          - leftNodeInfo.getQuote().length();
-        int quotePos = rightNodeInfo.getAbsoluteContentRange().getStartOffset() + replacement.length();
-        req.document.insertString(quotePos, rightNodeInfo.getQuote());
+        cutIntoRight = replacement.length() + rightElem.getQuote().length();
+        int contentWithQuoteStartColumn = leftElem.getTextOffset() + leftElem.getPrefixLength() - getLeftLineStartOffset(req);
+        int quotePos = rightElem.getTextOffset() + rightElem.getContentRange().getStartOffset() + replacement.length();
+        req.document.insertString(quotePos, rightElem.getQuote());
         req.document.insertString(quotePos, StringUtil.repeat(" ", contentWithQuoteStartColumn));
         lineEnd = quote + "\\\n";
       }
-      return new Result(replacement + lineEnd, 0, leftNodeInfo.getQuote().length(), cutIntoRight);
+      return new Result(replacement + lineEnd, 0, leftElem.getQuote().length(), cutIntoRight);
     }
 
-    private static boolean haveSamePrefixes(@NotNull StringNodeInfo leftNodeInfo, @NotNull StringNodeInfo rightNodeInfo) {
+    private static boolean haveSamePrefixes(@NotNull PyStringElement leftNodeInfo, @NotNull PyStringElement rightNodeInfo) {
       return leftNodeInfo.isUnicode() == rightNodeInfo.isUnicode() &&
              leftNodeInfo.isRaw() == rightNodeInfo.isRaw() &&
-             leftNodeInfo.isBytes() == rightNodeInfo.isBytes();
+             leftNodeInfo.isBytes() == rightNodeInfo.isBytes() &&
+             // TODO Merge formatted and plain strings elements (escape curly braces, etc.)
+             leftNodeInfo.isFormatted() == rightNodeInfo.isFormatted();
     }
   }
 
diff --git a/python/testData/joinLines/FStringAndPlainLiteral-after.py b/python/testData/joinLines/FStringAndPlainLiteral-after.py
new file mode 100644
index 000000000000..76353360971b
--- /dev/null
+++ b/python/testData/joinLines/FStringAndPlainLiteral-after.py
@@ -0,0 +1 @@
+(f'{42}' '{}')
\ No newline at end of file
diff --git a/python/testData/joinLines/FStringAndPlainLiteral.py b/python/testData/joinLines/FStringAndPlainLiteral.py
new file mode 100644
index 000000000000..c977b098e146
--- /dev/null
+++ b/python/testData/joinLines/FStringAndPlainLiteral.py
@@ -0,0 +1,2 @@
+(f'{42}'
+'{}')
\ No newline at end of file
diff --git a/python/testData/joinLines/FStrings-after.py b/python/testData/joinLines/FStrings-after.py
new file mode 100644
index 000000000000..0bbaeab309e8
--- /dev/null
+++ b/python/testData/joinLines/FStrings-after.py
@@ -0,0 +1 @@
+f'{42}foo'
diff --git a/python/testData/joinLines/FStrings.py b/python/testData/joinLines/FStrings.py
new file mode 100644
index 000000000000..3b37b48afbce
--- /dev/null
+++ b/python/testData/joinLines/FStrings.py
@@ -0,0 +1,2 @@
+f'{42}'
+f'foo'
diff --git a/python/testSrc/com/jetbrains/python/PyJoinLinesTest.java b/python/testSrc/com/jetbrains/python/PyJoinLinesTest.java
index 93ab8b674814..e3506cb91065 100644
--- a/python/testSrc/com/jetbrains/python/PyJoinLinesTest.java
+++ b/python/testSrc/com/jetbrains/python/PyJoinLinesTest.java
@@ -178,4 +178,12 @@ public class PyJoinLinesTest extends PyTestCase {
   public void testBackslashBetweenTargetsInFromImport() {
     doTest();
   }
+
+  public void testFStrings() {
+    doTest();
+  }
+  
+  public void testFStringAndPlainLiteral() {
+    doTest();
+  }
 }