From 7183fd55b56addc0dc4d279614bbeb7c7048559d Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Tue, 3 Mar 2020 19:09:04 +0300
Subject: [PATCH] PY-40778 Keep spaces between f-string fragment own braces and
 that of embedded expression

GitOrigin-RevId: f392f683b0c114fb046ede5b317bc4b0be207ec1
---
 .../com/jetbrains/python/formatter/PyBlock.java | 17 +++++++++++++++--
 ...tweenFragmentAndExpressionBracesPreserved.py |  8 ++++++++
 ...ragmentAndExpressionBracesPreserved_after.py |  8 ++++++++
 .../com/jetbrains/python/PyFormatterTest.java   |  5 +++++
 4 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 python/testData/formatter/fStringSpacesBetweenFragmentAndExpressionBracesPreserved.py
 create mode 100644 python/testData/formatter/fStringSpacesBetweenFragmentAndExpressionBracesPreserved_after.py

diff --git a/python/src/com/jetbrains/python/formatter/PyBlock.java b/python/src/com/jetbrains/python/formatter/PyBlock.java
index a6279fdd97e3..c2a5bc813d00 100644
--- a/python/src/com/jetbrains/python/formatter/PyBlock.java
+++ b/python/src/com/jetbrains/python/formatter/PyBlock.java
@@ -6,6 +6,7 @@ import com.intellij.lang.ASTNode;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
+import com.intellij.psi.impl.source.tree.LeafElement;
 import com.intellij.psi.impl.source.tree.TreeUtil;
 import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.tree.TokenSet;
@@ -821,6 +822,19 @@ public class PyBlock implements ASTBlock {
       if (isInsideFStringFragmentWithEqualsSign(myNode)) {
         return Spacing.getReadOnlySpacing();
       }
+      if (childType1 == PyTokenTypes.FSTRING_FRAGMENT_START) {
+        final LeafElement firstLeaf = TreeUtil.findFirstLeaf(node2);
+        if (firstLeaf != null && firstLeaf.getElementType() == PyTokenTypes.LBRACE) {
+          return Spacing.createSpacing(1, 1, 0, false, 0);
+        }
+      }
+
+      if (childType2 == PyTokenTypes.FSTRING_FRAGMENT_END) {
+        final ASTNode lastLeaf = TreeUtil.findLastLeaf(node1);
+        if (lastLeaf != null && lastLeaf.getElementType() == PyTokenTypes.RBRACE) {
+          return Spacing.createSpacing(1, 1, 0, false, 0);
+        }
+      }
 
       if ((childType1 == PyTokenTypes.EQ || childType2 == PyTokenTypes.EQ)) {
         final PyNamedParameter namedParameter = as(myNode.getPsi(), PyNamedParameter.class);
@@ -1176,8 +1190,7 @@ public class PyBlock implements ASTBlock {
   private static boolean isInsideFStringFragmentWithEqualsSign(@NotNull ASTNode node) {
     final ASTNode fStringFragmentParent = node.getElementType() == PyElementTypes.FSTRING_FRAGMENT
                                     ? node
-                                    : TreeUtil.findParent(node, TokenSet.create(PyElementTypes.FSTRING_FRAGMENT),
-                                                          stopAtTokens);
+                                    : TreeUtil.findParent(node, TokenSet.create(PyElementTypes.FSTRING_FRAGMENT), stopAtTokens);
     if (fStringFragmentParent == null) return false;
     return fStringFragmentParent.findChildByType(PyTokenTypes.EQ) != null;
   }
diff --git a/python/testData/formatter/fStringSpacesBetweenFragmentAndExpressionBracesPreserved.py b/python/testData/formatter/fStringSpacesBetweenFragmentAndExpressionBracesPreserved.py
new file mode 100644
index 000000000000..4f9578970487
--- /dev/null
+++ b/python/testData/formatter/fStringSpacesBetweenFragmentAndExpressionBracesPreserved.py
@@ -0,0 +1,8 @@
+s = f'{ {} }'
+s1 = f'{ {1} }'
+s2 = f'{ {1: 2} }'
+s3 = f'{ {x for x in "foo"} }'
+s4 = f'{ {x: 42 for x in "foo"} }'
+s5 = f'{ {1} + {2} }'
+s6 = f'{ {1} if True else {2} }'
+s7 = f'{ {1}.intersection({2}) }'
diff --git a/python/testData/formatter/fStringSpacesBetweenFragmentAndExpressionBracesPreserved_after.py b/python/testData/formatter/fStringSpacesBetweenFragmentAndExpressionBracesPreserved_after.py
new file mode 100644
index 000000000000..e04c258937d1
--- /dev/null
+++ b/python/testData/formatter/fStringSpacesBetweenFragmentAndExpressionBracesPreserved_after.py
@@ -0,0 +1,8 @@
+s = f'{ {} }'
+s1 = f'{ {1} }'
+s2 = f'{ {1: 2} }'
+s3 = f'{ {x for x in "foo"} }'
+s4 = f'{ {x: 42 for x in "foo"} }'
+s5 = f'{ {1} + {2} }'
+s6 = f'{ {1} if True else {2} }'
+s7 = f'{ {1}.intersection({2})}'
diff --git a/python/testSrc/com/jetbrains/python/PyFormatterTest.java b/python/testSrc/com/jetbrains/python/PyFormatterTest.java
index 71529c2d5375..122b975da663 100644
--- a/python/testSrc/com/jetbrains/python/PyFormatterTest.java
+++ b/python/testSrc/com/jetbrains/python/PyFormatterTest.java
@@ -956,6 +956,11 @@ public class PyFormatterTest extends PyTestCase {
     runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
   }
 
+  // PY-40778
+  public void testFStringSpacesBetweenFragmentAndExpressionBracesPreserved() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doTest);
+  }
+
   // PY-31991
   public void testSpacesAroundFStringFragmentExpressionStripped() {
     runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);