PY-12018 First attempt to sort imports alphabetically in ImportOptimizer

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

@@ -15,11 +15,14 @@
  */
 package com.jetbrains.python.codeInsight.imports;
 
+import com.google.common.base.Function;
+import com.google.common.collect.Ordering;
 import com.intellij.codeInspection.LocalInspectionToolSession;
 import com.intellij.lang.ImportOptimizer;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiFileSystemItem;
+import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.ArrayUtil;
 import com.jetbrains.python.formatter.PyBlock;
 import com.jetbrains.python.inspections.unresolvedReference.PyUnresolvedReferencesInspection;
@@ -59,13 +62,33 @@ public class PyImportOptimizer implements ImportOptimizer {
       public void run() {
         visitor.optimizeImports();
         if (file instanceof PyFile) {
-          new ImportSorter((PyFile) file).run();
+          new ImportSorter((PyFile)file).run();
         }
       }
     };
   }
 
   private static class ImportSorter {
+    public static final Ordering<PyImportStatementBase> BY_NAME_COMPARATOR =
+      Ordering.natural()
+              .nullsLast()
+              .onResultOf(new Function<PyImportStatementBase, String>() {
+                @Override
+                public String apply(@Nullable PyImportStatementBase importStatement) {
+                  QualifiedName qualifiedName = null;
+                  if (importStatement instanceof PyFromImportStatement) {
+                    qualifiedName = ((PyFromImportStatement)importStatement).getImportSourceQName();
+                  }
+                  else if (importStatement instanceof PyImportStatement) {
+                    final PyImportElement importElement = ArrayUtil.getFirstElement(importStatement.getImportElements());
+                    if (importElement != null) {
+                      qualifiedName = importElement.getImportedQName();
+                    }
+                  }
+                  return qualifiedName == null ? null : qualifiedName.toString();
+                }
+              });
+
     private final PyFile myFile;
     private final List<PyImportStatementBase> myBuiltinImports = new ArrayList<PyImportStatementBase>();
     private final List<PyImportStatementBase> myThirdPartyImports = new ArrayList<PyImportStatementBase>();
@@ -92,19 +115,19 @@ public class PyImportOptimizer implements ImportOptimizer {
         if (importStatement instanceof PyImportStatement && importStatement.getImportElements().length > 1) {
           for (PyImportElement importElement : importStatement.getImportElements()) {
             myMissorted = true;
-            PsiElement toImport = importElement.resolve();
+            // getText() for ImportElement includes alias
             final PyImportStatement splitImport = myGenerator.createImportStatement(langLevel, importElement.getText(), null);
-            prioritize(splitImport, toImport);
+            prioritize(splitImport, importElement.resolve());
           }
         }
         else {
-          PsiElement toImport;
+          final PsiElement toImport;
           if (importStatement instanceof PyFromImportStatement) {
             toImport = ((PyFromImportStatement)importStatement).resolveImportSource();
           }
           else {
             final PyImportElement firstImportElement = ArrayUtil.getFirstElement(importStatement.getImportElements());
-            toImport = firstImportElement != null? firstImportElement.resolve() : null;
+            toImport = firstImportElement != null ? firstImportElement.resolve() : null;
           }
           prioritize(importStatement, toImport);
         }
@@ -147,13 +170,18 @@ public class PyImportOptimizer implements ImportOptimizer {
     }
 
     private void applyResults() {
+      Collections.sort(myBuiltinImports, BY_NAME_COMPARATOR);
+      Collections.sort(myThirdPartyImports, BY_NAME_COMPARATOR);
+      Collections.sort(myProjectImports, BY_NAME_COMPARATOR);
+
       markGroupBegin(myThirdPartyImports);
       markGroupBegin(myProjectImports);
+
       addImports(myBuiltinImports);
       addImports(myThirdPartyImports);
       addImports(myProjectImports);
-      PsiElement lastElement = myImportBlock.get(myImportBlock.size()-1);
-      PyImportStatementBase firstNonFutureImport = findFirstNonFutureImport();
+      final PsiElement lastElement = myImportBlock.get(myImportBlock.size() - 1);
+      final PyImportStatementBase firstNonFutureImport = findFirstNonFutureImport();
       if (firstNonFutureImport != null) {
         myFile.deleteChildRange(firstNonFutureImport, lastElement);
       }
@@ -163,7 +191,7 @@ public class PyImportOptimizer implements ImportOptimizer {
     }
 
     private PyImportStatementBase findFirstNonFutureImport() {
-      for (PyImportStatementBase importStatement: myImportBlock) {
+      for (PyImportStatementBase importStatement : myImportBlock) {
         if (!(importStatement instanceof PyFromImportStatement && ((PyFromImportStatement)importStatement).isFromFuture())) {
           return importStatement;
         }
@@ -171,14 +199,14 @@ public class PyImportOptimizer implements ImportOptimizer {
       return null;
     }
 
-    private static void markGroupBegin(List<PyImportStatementBase> imports) {
+    private static void markGroupBegin(@NotNull List<PyImportStatementBase> imports) {
       if (imports.size() > 0) {
         imports.get(0).putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, true);
       }
     }
 
     private void addImports(final List<PyImportStatementBase> imports) {
-      for (PyImportStatementBase newImport: imports) {
+      for (PyImportStatementBase newImport : imports) {
         myFile.addBefore(newImport, findFirstNonFutureImport());
       }
     }

python/testData/optimizeImports/alphabeticalOrder.after.py

@@ -0,0 +1,13 @@
+from datetime import timedelta
+import sys
+
+import a
+from a import C1
+import b
+from b import func
+import foo # broken
+import z
+from
+import # broken
+
+print(z, b, a, C1, func, sys, abc, foo, timedelta)

python/testData/optimizeImports/alphabeticalOrder.py

@@ -0,0 +1,11 @@
+import z
+import b
+import a
+from a import C1
+from b import func
+from import foo # broken
+import sys
+from datetime import timedelta
+import # broken
+
+print(z, b, a, C1, func, sys, abc, foo, timedelta)

python/testData/optimizeImports/insertBlankLines.after.py

@@ -1,9 +1,9 @@
 from __future__ import with_statement
-import sys
 import datetime
+import sys
 
-import foo
 from bar import *
+import foo
 
 sys.path
 datetime.datetime

python/testData/optimizeImports/orderByType.after.py

@@ -1,9 +1,9 @@
 from __future__ import with_statement
-import sys
 import datetime
+import sys
 
-import foo
 from bar import *
+import foo
 
 sys.path
 datetime.datetime

python/testData/optimizeImports/split.after.py

@@ -1,5 +1,5 @@
-import sys
 import datetime
+import sys
 
 sys.path
 datetime.time

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

@@ -59,7 +59,12 @@ public class PyOptimizeImportsTest extends PyTestCase {
     doTest();
   }
 
-  public void testOrder() {
+  public void testOrderByType() {
+    doTest();
+  }
+
+  // PY-12018
+  public void testAlphabeticalOrder() {
     doTest();
   }