From c5fa661c0f766aaa171a4b59ae398a5a0e7a73e6 Mon Sep 17 00:00:00 2001
From: Ekaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Date: Mon, 21 Jan 2013 14:01:35 +0400
Subject: [PATCH] improved reflow paragraph action

---
 .../actions/RestFillParagraphHandler.java     | 15 +++++++++++++
 .../actions/PyFillParagraphHandler.java       | 22 +++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/python/rest/src/com/jetbrains/rest/actions/RestFillParagraphHandler.java b/python/rest/src/com/jetbrains/rest/actions/RestFillParagraphHandler.java
index f5665e349248..476f176945ad 100644
--- a/python/rest/src/com/jetbrains/rest/actions/RestFillParagraphHandler.java
+++ b/python/rest/src/com/jetbrains/rest/actions/RestFillParagraphHandler.java
@@ -20,6 +20,12 @@ public class RestFillParagraphHandler extends ParagraphFillHandler {
     return element instanceof PsiComment? ".. " : "";
   }
 
+  @NotNull
+  @Override
+  protected String getPostfix(@NotNull PsiElement element) {
+    return element.getNode().getElementType() == RestTokenTypes.COMMENT? "\n" : "";
+  }
+
   @Override
   protected boolean isAvailableForFile(@Nullable PsiFile psiFile) {
     return psiFile instanceof RestFile;
@@ -35,4 +41,13 @@ public class RestFillParagraphHandler extends ParagraphFillHandler {
     return element instanceof PsiWhiteSpace ||
            element != null && element.getNode().getElementType() == RestTokenTypes.WHITESPACE;
   }
+
+  @Override
+  protected void appendPostfix(@NotNull PsiElement element,
+                               @NotNull String text,
+                               @NotNull StringBuilder stringBuilder) {
+    if (element.getNode().getElementType() == RestTokenTypes.COMMENT) {
+      stringBuilder.append(getPostfix(element));
+    }
+  }
 }
diff --git a/python/src/com/jetbrains/python/actions/PyFillParagraphHandler.java b/python/src/com/jetbrains/python/actions/PyFillParagraphHandler.java
index 281a97494081..51e04707135d 100644
--- a/python/src/com/jetbrains/python/actions/PyFillParagraphHandler.java
+++ b/python/src/com/jetbrains/python/actions/PyFillParagraphHandler.java
@@ -1,10 +1,12 @@
 package com.jetbrains.python.actions;
 
 import com.intellij.codeInsight.editorActions.fillParagraph.ParagraphFillHandler;
+import com.intellij.openapi.util.Pair;
 import com.intellij.psi.PsiComment;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.PythonStringUtil;
 import com.jetbrains.python.psi.PyFile;
 import com.jetbrains.python.psi.PyStringLiteralExpression;
 import org.jetbrains.annotations.NotNull;
@@ -17,9 +19,29 @@ public class PyFillParagraphHandler extends ParagraphFillHandler {
 
   @NotNull
   protected String getPrefix(@NotNull final PsiElement element) {
+    final PyStringLiteralExpression stringLiteralExpression =
+      PsiTreeUtil.getParentOfType(element, PyStringLiteralExpression.class);
+    if (stringLiteralExpression != null) {
+      final Pair<String,String> quotes =
+        PythonStringUtil.getQuotes(stringLiteralExpression.getText());
+      return quotes != null? quotes.getFirst()+"\n" : "\"\n";
+    }
     return element instanceof PsiComment? "#" : "";
   }
 
+  @NotNull
+  @Override
+  protected String getPostfix(@NotNull PsiElement element) {
+    final PyStringLiteralExpression stringLiteralExpression =
+      PsiTreeUtil.getParentOfType(element, PyStringLiteralExpression.class);
+    if (stringLiteralExpression != null) {
+      final Pair<String,String> quotes =
+        PythonStringUtil.getQuotes(stringLiteralExpression.getText());
+      return quotes != null? "\n" + quotes.getSecond() : "\n\"";
+    }
+    return "";
+  }
+
   @Override
   protected boolean isAvailableForElement(@Nullable PsiElement element) {
     if (element != null) {