PY-22355 Optimize imports honors new code style settings for "from" imports

Namely, it automatically surrounds them with parentheses and inserts
a trailing comma after the last imported name if necessary, e.g. when
several "from" imports with the same source were combined into a
single long statement that needs to be wrapped.

python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java

@@ -22,6 +22,7 @@ import com.intellij.openapi.util.Comparing;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
 import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.containers.MultiMap;
@@ -226,12 +227,25 @@ public class PyImportOptimizer implements ImportOptimizer {
           myGroups.put(priority, imports);
         }
       }
+      prepareNewImports();
       markGroupStarts();
       addImports(myImportBlock.get(0));
 
       myFile.deleteChildRange(myImportBlock.get(0), myImportBlock.get(myImportBlock.size() - 1));
     }
 
+    private void prepareNewImports() {
+      for (List<PyImportStatementBase> imports : myGroups.values()) {
+        for (int i = 0; i < imports.size(); i++) {
+          final PyImportStatementBase newImport = imports.get(i);
+          final CodeStyleManager styleManager = CodeStyleManager.getInstance(newImport.getProject());
+          // Some of imports were copied as is and they're still present in the original PSI file
+          final PyImportStatementBase formatted = (PyImportStatementBase)styleManager.reformat(newImport.copy());
+          imports.set(i, formatted);
+        }
+      }
+    }
+
     private void markGroupStarts() {
       for (List<PyImportStatementBase> group : myGroups.values()) {
         boolean firstImportInGroup = true;

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

@@ -16,10 +16,7 @@
 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.*;
 import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.codeStyle.CodeStyleSettings;
 import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
@@ -27,6 +24,7 @@ import com.intellij.psi.impl.source.codeStyle.PostFormatProcessor;
 import com.intellij.psi.impl.source.codeStyle.PostFormatProcessorHelper;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyTokenTypes;
+import com.jetbrains.python.PythonLanguage;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import org.jetbrains.annotations.NotNull;
@@ -53,6 +51,7 @@ public class PyFromImportPostFormatProcessor implements PostFormatProcessor {
   private static class Visitor extends PyRecursiveElementVisitor {
     private final PostFormatProcessorHelper myHelper;
     private final List<PyFromImportStatement> myImportStatements = new ArrayList<>();
+    private PsiElement myRootElement;
 
     public Visitor(@NotNull CommonCodeStyleSettings settings) {
       myHelper = new PostFormatProcessorHelper(settings);
@@ -83,8 +82,10 @@ public class PyFromImportPostFormatProcessor implements PostFormatProcessor {
 
     @NotNull
     public PsiElement processElement(@NotNull PsiElement element) {
+      // For some reason smart pointers don't work for non-physical (in particular, generated) elements
+      myRootElement = element;
       findAndReplaceFromImports(element);
-      return element;
+      return myRootElement;
     }
 
     @NotNull
@@ -95,11 +96,15 @@ public class PyFromImportPostFormatProcessor implements PostFormatProcessor {
     }
 
     private void findAndReplaceFromImports(@NotNull PsiElement element) {
-      if (element.getContainingFile() instanceof PyFile) {
+      // Copied/generated elements are stored in DummyHolder files, not PyFiles
+      if (element.getLanguage().isKindOf(PythonLanguage.INSTANCE)) {
         element.accept(this);
         Collections.reverse(myImportStatements);
         for (PyFromImportStatement statement : myImportStatements) {
-          replaceFromImport(statement);
+          final PyFromImportStatement newStatement = replaceFromImport(statement);
+          if (myRootElement == statement) {
+            myRootElement = newStatement;
+          }
         }
       }
     }

python/testData/optimizeImports/alphabeticalOrder.after.py

@@ -4,10 +4,10 @@ from __future__ import unicode_literals
 import sys
 from datetime import timedelta
 
-import # broken
+import  # broken
 import a
 import b
-import foo # broken
+import foo  # broken
 import z
 from
 from a import C1

python/testData/optimizeImports/disableAlphabeticalOrder.after.py

@@ -11,8 +11,8 @@ from a import C1
 from alphabet import D
 from b import func
 from
-import foo # broken
-import # broken
+import foo  # broken
+import  # broken
 from alphabet import *
 from alphabet import C
 from alphabet import B, A

python/testData/optimizeImports/parenthesesAndTrailingCommaInFromImports.after.py

@@ -0,0 +1,6 @@
+from statistics import (
+    median,
+    variance,
+)
+
+print(   median,variance)#poorly formatted, but shouldn't be affected

python/testData/optimizeImports/parenthesesAndTrailingCommaInFromImports.py

@@ -0,0 +1,4 @@
+from statistics import median
+from statistics import variance
+
+print(   median,variance)#poorly formatted, but shouldn't be affected

python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java

@@ -22,6 +22,7 @@ import com.intellij.openapi.projectRoots.Sdk;
 import com.intellij.openapi.projectRoots.SdkModificator;
 import com.intellij.openapi.roots.OrderRootType;
 import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
 import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.PyImportStatementBase;
@@ -230,6 +231,17 @@ public class PyOptimizeImportsTest extends PyTestCase {
     runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest);
   }
 
+  // PY-22355
+  public void testParenthesesAndTrailingCommaInFromImports() {
+    getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_JOIN_FROM_IMPORTS_WITH_SAME_SOURCE = true;
+    getPythonCodeStyleSettings().FROM_IMPORT_WRAPPING = CommonCodeStyleSettings.WRAP_ALWAYS;
+    getPythonCodeStyleSettings().FROM_IMPORT_PARENTHESES_FORCE_IF_MULTILINE = true;
+    getPythonCodeStyleSettings().FROM_IMPORT_NEW_LINE_AFTER_LEFT_PARENTHESIS = true;
+    getPythonCodeStyleSettings().FROM_IMPORT_NEW_LINE_BEFORE_RIGHT_PARENTHESIS = true;
+    getPythonCodeStyleSettings().FROM_IMPORT_TRAILING_COMMA_IF_MULTILINE = true;
+    doTest();
+  }
+
   private void doTest() {
     myFixture.configureByFile(getTestName(true) + ".py");
     OptimizeImportsAction.actionPerformedImpl(DataManager.getInstance().getDataContext(myFixture.getEditor().getContentComponent()));