PY-20633 Allow to forcibly surround "from" imports with parentheses

As in case with similar options for braces in Java and other languages
there are three variants of behavior: don't force, force always, and
force if the corresponding statement spans multiple lines.
In either way, the option is applied only if "from" import contains
more than one imported name.

PY-20633 Allow to wrap only multiline "from" imports in parentheses

python/src/META-INF/python-core-common.xml

@@ -28,6 +28,7 @@
     <lang.formatter language="Python" implementationClass="com.jetbrains.python.formatter.PythonFormattingModelBuilder"/>
     <preFormatProcessor implementation="com.jetbrains.python.formatter.PyPreFormatProcessor"/>
     <postFormatProcessor implementation="com.jetbrains.python.formatter.PyTrailingBlankLinesPostFormatProcessor"/>
+    <postFormatProcessor implementation="com.jetbrains.python.formatter.PyFromImportPostFormatProcessor"/>
     <lang.whiteSpaceFormattingStrategy language="Python"
                                        implementationClass="com.jetbrains.python.formatter.PyWhiteSpaceFormattingStrategy"/>
     <lang.lineWrapStrategy language="Python" implementationClass="com.jetbrains.python.formatter.PyLineWrapPositionStrategy"/>

python/src/com/jetbrains/python/PyBundle.properties

@@ -1034,6 +1034,7 @@ formatter.force.new.line.after.colon=Force new line after colon
 formatter.align.when.multiline=Align when multiline
 formatter.collections.and.comprehensions=Collections and Comprehensions
 formatter.from.import.statements="From" Import Statements
+formatter.from.import.statements.force.parentheses=Force parentheses
 formatter.dictionary.literals=Dictionary literals
 formatter.hang.closing.brackets=Hang closing brackets
 

python/src/com/jetbrains/python/formatter/PyCodeStyleSettings.java

@@ -17,8 +17,10 @@ package com.jetbrains.python.formatter;
 
 import com.intellij.formatting.WrapType;
 import com.intellij.psi.codeStyle.CodeStyleSettings;
+import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
 import com.intellij.psi.codeStyle.CustomCodeStyleSettings;
 import com.jetbrains.python.PyBundle;
+import org.intellij.lang.annotations.MagicConstant;
 
 /**
  * @author yole
@@ -74,6 +76,12 @@ public class PyCodeStyleSettings extends CustomCodeStyleSettings {
   public boolean SPACE_BEFORE_NUMBER_SIGN = true;
 
   public int DICT_ALIGNMENT = DICT_ALIGNMENT_NONE;
+  @MagicConstant(intValues = {
+    CommonCodeStyleSettings.DO_NOT_WRAP,
+    CommonCodeStyleSettings.WRAP_AS_NEEDED,
+    CommonCodeStyleSettings.WRAP_ALWAYS,
+    CommonCodeStyleSettings.WRAP_ON_EVERY_ITEM
+  })
   public int DICT_WRAPPING = WrapType.NORMAL.getLegacyRepresentation();
   public boolean DICT_NEW_LINE_AFTER_LEFT_BRACE = false;
   public boolean DICT_NEW_LINE_BEFORE_RIGHT_BRACE = false;
@@ -96,9 +104,22 @@ public class PyCodeStyleSettings extends CustomCodeStyleSettings {
   /**
    * Affects wrapping of multiple imported names in a single "from" import.
    */
+  @MagicConstant(intValues = {
+    CommonCodeStyleSettings.DO_NOT_WRAP,
+    CommonCodeStyleSettings.WRAP_AS_NEEDED,
+    CommonCodeStyleSettings.WRAP_ALWAYS,
+    CommonCodeStyleSettings.WRAP_ON_EVERY_ITEM
+  })
   public int FROM_IMPORT_WRAPPING = WrapType.NORMAL.getLegacyRepresentation();
   public boolean FROM_IMPORT_NEW_LINE_AFTER_LEFT_PARENTHESIS = false;
   public boolean FROM_IMPORT_NEW_LINE_BEFORE_RIGHT_PARENTHESIS = false;
+  
+  @MagicConstant(intValues = {
+    CommonCodeStyleSettings.DO_NOT_FORCE,
+    CommonCodeStyleSettings.FORCE_BRACES_IF_MULTILINE,
+    CommonCodeStyleSettings.FORCE_BRACES_ALWAYS}
+  )
+  public int FROM_IMPORT_PARENTHESES_FORCE = CommonCodeStyleSettings.DO_NOT_FORCE;
 
   /**
    * Corresponds to the option of pycodestyle.py "--hang-closing". Basically, it means that the closing brace of a collection literal, 

python/src/com/jetbrains/python/formatter/PyFromImportPostFormatProcessor.java

@@ -0,0 +1,128 @@
+/*
+ * Copyright 2000-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.formatter;
+
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiComment;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiWhiteSpace;
+import com.intellij.psi.codeStyle.CodeStyleManager;
+import com.intellij.psi.codeStyle.CodeStyleSettings;
+import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
+import com.intellij.psi.impl.source.codeStyle.PostFormatProcessor;
+import com.intellij.psi.impl.source.codeStyle.PostFormatProcessorHelper;
+import com.jetbrains.python.psi.*;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * @author Mikhail Golubev
+ */
+public class PyFromImportPostFormatProcessor implements PostFormatProcessor {
+  @Override
+  public PsiElement processElement(@NotNull PsiElement source, @NotNull CodeStyleSettings settings) {
+    return new Visitor(settings).processElement(source);
+  }
+
+  @Override
+  public TextRange processText(@NotNull PsiFile source, @NotNull TextRange rangeToReformat, @NotNull CodeStyleSettings settings) {
+    return new Visitor(settings).processTextRange(source, rangeToReformat);
+  }
+
+  private static class Visitor extends PyRecursiveElementVisitor {
+    private final PostFormatProcessorHelper myHelper;
+    private final List<PyFromImportStatement> myImportStatements = new ArrayList<>();
+
+    public Visitor(@NotNull CommonCodeStyleSettings settings) {
+      myHelper = new PostFormatProcessorHelper(settings);
+    }
+
+    @Override
+    public void visitPyFromImportStatement(PyFromImportStatement node) {
+      if (myHelper.isElementFullyInRange(node)) {
+        final PyImportElement[] importedNames = node.getImportElements();
+        final PyCodeStyleSettings pySettings = ((CodeStyleSettings)myHelper.getSettings()).getCustomSettings(PyCodeStyleSettings.class);
+        final PsiElement leftParen = node.getLeftParen();
+        final boolean enabledInSettings = pySettings.FROM_IMPORT_PARENTHESES_FORCE == CommonCodeStyleSettings.FORCE_BRACES_ALWAYS ||
+                                          pySettings.FROM_IMPORT_PARENTHESES_FORCE == CommonCodeStyleSettings.FORCE_BRACES_IF_MULTILINE &&
+                                          PostFormatProcessorHelper.isMultiline(node);
+        if (enabledInSettings && importedNames.length > 1 && leftParen == null) {
+          myImportStatements.add(node);
+        }
+      }
+    }
+
+    @NotNull
+    public PsiElement processElement(@NotNull PsiElement element) {
+      findAndReplaceFromImports(element);
+      return element;
+    }
+
+    @NotNull
+    public TextRange processTextRange(@NotNull PsiFile file, @NotNull TextRange range) {
+      myHelper.setResultTextRange(range);
+      findAndReplaceFromImports(file);
+      return myHelper.getResultTextRange();
+    }
+
+    private void findAndReplaceFromImports(@NotNull PsiElement element) {
+      if (element.getContainingFile() instanceof PyFile) {
+        element.accept(this);
+        Collections.reverse(myImportStatements);
+        for (PyFromImportStatement statement : myImportStatements) {
+          replaceFromImport(statement);
+        }
+      }
+    }
+
+    @NotNull
+    private PyFromImportStatement replaceFromImport(@NotNull PyFromImportStatement fromImport) {
+      final PyImportElement firstName = fromImport.getImportElements()[0];
+      final String beforeFirstName = fromImport.getText().substring(0, firstName.getStartOffsetInParent());
+      final StringBuilder newStatementText = new StringBuilder(beforeFirstName);
+      newStatementText.append("(");
+      boolean lastElementWasComment = false;
+      for (PsiElement cur = firstName; cur != null; cur = cur.getNextSibling()) {
+        if (cur instanceof PsiWhiteSpace) {
+          newStatementText.append(cur.getText().replace("\\", ""));
+        }
+        else {
+          newStatementText.append(cur.getText());
+        }
+        lastElementWasComment = cur instanceof PsiComment;
+      }
+      if (lastElementWasComment) {
+        newStatementText.append("\n");
+      }
+      newStatementText.append(")");
+      
+      final PyElementGenerator generator = PyElementGenerator.getInstance(fromImport.getProject());
+      final CodeStyleManager codeStyleManager = CodeStyleManager.getInstance(fromImport.getProject());
+      
+      final LanguageLevel level = LanguageLevel.forElement(fromImport);
+      PyFromImportStatement newFromImport = generator.createFromText(level, PyFromImportStatement.class, newStatementText.toString());
+      newFromImport = (PyFromImportStatement)fromImport.replace(newFromImport);
+      newFromImport = (PyFromImportStatement)codeStyleManager.reformat(newFromImport);
+      myHelper.updateResultRange(fromImport.getTextLength(), newFromImport.getTextLength());
+      return newFromImport;
+    }
+  }
+  
+}

python/src/com/jetbrains/python/formatter/PyLanguageCodeStyleSettingsProvider.java

@@ -131,6 +131,10 @@ public class PyLanguageCodeStyleSettingsProvider extends LanguageCodeStyleSettin
       consumer.showCustomOption(PyCodeStyleSettings.class, "FROM_IMPORT_NEW_LINE_BEFORE_RIGHT_PARENTHESIS",
                                 ApplicationBundle.message("wrapping.rpar.on.new.line"),
                                 PyBundle.message("formatter.from.import.statements"));
+      consumer.showCustomOption(PyCodeStyleSettings.class, "FROM_IMPORT_PARENTHESES_FORCE",
+                                PyBundle.message("formatter.from.import.statements.force.parentheses"),
+                                PyBundle.message("formatter.from.import.statements"),
+                                BRACE_OPTIONS, BRACE_VALUES);
 
       consumer.showCustomOption(PyCodeStyleSettings.class, "DICT_WRAPPING",
                                 PyBundle.message("formatter.dictionary.literals"), null, WRAP_OPTIONS, WRAP_VALUES);

python/testData/formatter/fromImportForceParenthesesAlways.py

@@ -0,0 +1,12 @@
+from module import foo
+from module import foo, bar
+#                            | margin
+from module import foo, bar, baz
+from module import foo, \
+    bar
+from module import foo, \
+    bar  # comment
+from module import foo, \
+    \
+    \
+    bar,

python/testData/formatter/fromImportForceParenthesesAlways_after.py

@@ -0,0 +1,13 @@
+from module import foo
+from module import (foo, bar)
+#                            | margin
+from module import (foo, bar,
+                    baz)
+from module import (foo,
+                    bar)
+from module import (foo,
+                    bar  # comment
+                    )
+from module import (foo,
+
+                    bar, )

python/testData/formatter/fromImportForceParenthesesIfMultiline.py

@@ -0,0 +1,12 @@
+from module import foo
+from module import foo, bar
+#                            | margin
+from module import foo, bar, baz
+from module import foo, \
+    bar
+from module import foo, \
+    bar  # comment
+from module import foo, \
+    \
+    \
+    bar,

python/testData/formatter/fromImportForceParenthesesIfMultiline_after.py

@@ -0,0 +1,13 @@
+from module import foo
+from module import foo, bar
+#                            | margin
+from module import (foo, bar,
+                    baz)
+from module import (foo,
+                    bar)
+from module import (foo,
+                    bar  # comment
+                    )
+from module import (foo,
+
+                    bar, )

python/testData/formatter/pyramidFromImportFormatting.py

@@ -0,0 +1,10 @@
+from module import foo
+from module import foo, bar
+from module import foo, \
+    bar
+from module import foo, \
+    bar  # comment
+from module import foo, \
+    \
+    \
+    bar,

python/testData/formatter/pyramidFromImportFormatting_after.py

@@ -0,0 +1,18 @@
+from module import foo
+from module import (
+    foo,
+    bar
+    )
+from module import (
+    foo,
+    bar
+    )
+from module import (
+    foo,
+    bar  # comment
+    )
+from module import (
+    foo,
+
+    bar,
+    )

python/testSrc/com/jetbrains/python/PyFormatterTest.java

@@ -20,6 +20,7 @@ import com.intellij.openapi.command.WriteCommandAction;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.codeStyle.CodeStyleManager;
+import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.formatter.PyCodeStyleSettings;
@@ -764,6 +765,31 @@ public class PyFormatterTest extends PyTestCase {
     doTest();
   }
 
+  // PY-20633
+  public void testFromImportForceParenthesesAlways() {
+    getCodeStyleSettings().setRightMargin(PythonLanguage.INSTANCE, 30);
+    getPythonCodeStyleSettings().FROM_IMPORT_PARENTHESES_FORCE = CommonCodeStyleSettings.FORCE_BRACES_ALWAYS;
+    doTest();
+  }
+
+  // PY-20633
+  public void testFromImportForceParenthesesIfMultiline() {
+    getCodeStyleSettings().setRightMargin(PythonLanguage.INSTANCE, 30);
+    getPythonCodeStyleSettings().FROM_IMPORT_PARENTHESES_FORCE = CommonCodeStyleSettings.FORCE_BRACES_IF_MULTILINE;
+    doTest();
+  }
+
+  // PY-20633
+  // See http://docs.pylonsproject.org/en/latest/community/codestyle.html
+  public void testPyramidFromImportFormatting() {
+    getPythonCodeStyleSettings().FROM_IMPORT_PARENTHESES_FORCE = CommonCodeStyleSettings.FORCE_BRACES_ALWAYS;
+    getPythonCodeStyleSettings().FROM_IMPORT_NEW_LINE_AFTER_LEFT_PARENTHESIS = true;
+    getPythonCodeStyleSettings().FROM_IMPORT_NEW_LINE_BEFORE_RIGHT_PARENTHESIS = true;
+    getPythonCodeStyleSettings().FROM_IMPORT_WRAPPING = WrapType.ALWAYS.getLegacyRepresentation();
+    getPythonCodeStyleSettings().HANG_CLOSING_BRACKETS = true;
+    doTest();
+  }
+
   public void testVariableAnnotations() {
     runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
   }