From d96766df2d207f25ca83c8302f8933b8eed41106 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Fri, 29 Aug 2014 16:10:09 +0400
Subject: [PATCH] PyFormatterModelBuilder always uses PyFile as root block
 element

The same way FormatterModelBuilders are implemented for Java and JavaScript.
It allows to call codeStyleManager.reformat(PsiElement) directly instead of
codeStyleManager.reformatText(elt.getContainingFile(),
elt.getTextRange().getStartOffset(), elt.getTextRange().getEndOffset());
---
 .../PythonFormattingModelBuilder.java         |  4 ++--
 .../reformatOfSingleElementPossible.py        |  2 ++
 .../reformatOfSingleElementPossible_after.py  |  2 ++
 .../com/jetbrains/python/PyFormatterTest.java | 22 +++++++++++++++++++
 4 files changed, 28 insertions(+), 2 deletions(-)
 create mode 100644 python/testData/formatter/reformatOfSingleElementPossible.py
 create mode 100644 python/testData/formatter/reformatOfSingleElementPossible_after.py

diff --git a/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java b/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java
index aee9e53adff6..4b71d6abec25 100644
--- a/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java
+++ b/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java
@@ -46,13 +46,13 @@ public class PythonFormattingModelBuilder implements FormattingModelBuilderEx, C
   public FormattingModel createModel(@NotNull PsiElement element,
                                      @NotNull CodeStyleSettings settings,
                                      @NotNull FormattingMode mode) {
+    final ASTNode fileNode = element.getContainingFile().getNode();
     if (DUMP_FORMATTING_AST) {
-      ASTNode fileNode = element.getContainingFile().getNode();
       System.out.println("AST tree for " + element.getContainingFile().getName() + ":");
       printAST(fileNode, 0);
     }
     final PyBlockContext context = new PyBlockContext(settings, createSpacingBuilder(settings), mode);
-    final PyBlock block = new PyBlock(null, element.getNode(), null, Indent.getNoneIndent(), null, context);
+    final PyBlock block = new PyBlock(null, fileNode, null, Indent.getNoneIndent(), null, context);
     if (DUMP_FORMATTING_AST) {
       FormattingModelDumper.dumpFormattingModel(block, 2, System.out);
     }
diff --git a/python/testData/formatter/reformatOfSingleElementPossible.py b/python/testData/formatter/reformatOfSingleElementPossible.py
new file mode 100644
index 000000000000..d97d8ab38f38
--- /dev/null
+++ b/python/testData/formatter/reformatOfSingleElementPossible.py
@@ -0,0 +1,2 @@
+x=[1,2,3]
+y='spam<caret>'*2
\ No newline at end of file
diff --git a/python/testData/formatter/reformatOfSingleElementPossible_after.py b/python/testData/formatter/reformatOfSingleElementPossible_after.py
new file mode 100644
index 000000000000..2528c032e2f1
--- /dev/null
+++ b/python/testData/formatter/reformatOfSingleElementPossible_after.py
@@ -0,0 +1,2 @@
+x=[1,2,3]
+y = 'spam' * 2
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyFormatterTest.java b/python/testSrc/com/jetbrains/python/PyFormatterTest.java
index ca0520c7b6a5..35a97e03efae 100644
--- a/python/testSrc/com/jetbrains/python/PyFormatterTest.java
+++ b/python/testSrc/com/jetbrains/python/PyFormatterTest.java
@@ -21,10 +21,12 @@ import com.intellij.psi.PsiFile;
 import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.codeStyle.CodeStyleSettings;
 import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.formatter.PyCodeStyleSettings;
 import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.PyElementGenerator;
+import com.jetbrains.python.psi.PyStatement;
 import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher;
 
 /**
@@ -424,6 +426,26 @@ public class PyFormatterTest extends PyTestCase {
     myFixture.checkResultByFile("formatter/" + getTestName(true) + "_after.py");
   }
 
+  /**
+   * This test merely checks that call to {@link com.intellij.psi.codeStyle.CodeStyleManager#reformat(com.intellij.psi.PsiElement)}
+   * is possible for Python sources.
+   */
+  public void testReformatOfSingleElementPossible() {
+    myFixture.configureByFile("formatter/" + getTestName(true) + ".py");
+    WriteCommandAction.runWriteCommandAction(myFixture.getProject(), new Runnable() {
+      @Override
+      public void run() {
+        final PsiElement elementAtCaret = myFixture.getFile().findElementAt(myFixture.getCaretOffset());
+        assertNotNull(elementAtCaret);
+        final PyStatement statement = PsiTreeUtil.getParentOfType(elementAtCaret, PyStatement.class, false);
+        assertNotNull(statement);
+        final CodeStyleManager codeStyleManager = CodeStyleManager.getInstance(myFixture.getProject());
+        codeStyleManager.reformat(statement);
+      }
+    });
+    myFixture.checkResultByFile("formatter/" + getTestName(true) + "_after.py");
+  }
+
   private CodeStyleSettings settings() {
     return CodeStyleSettingsManager.getInstance().getSettings(myFixture.getProject());
   }