import sorter skips from __future__ imports (PY-10022)

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

@@ -85,6 +85,9 @@ public class PyImportOptimizer implements ImportOptimizer {
       }
       LanguageLevel langLevel = LanguageLevel.forElement(myFile);
       for (PyImportStatementBase importStatement : myImportBlock) {
+        if (importStatement instanceof PyFromImportStatement && ((PyFromImportStatement)importStatement).isFromFuture()) {
+          continue;
+        }
         if (importStatement instanceof PyImportStatement && importStatement.getImportElements().length > 1) {
           for (PyImportElement importElement : importStatement.getImportElements()) {
             myMissorted = true;
@@ -96,7 +99,7 @@ public class PyImportOptimizer implements ImportOptimizer {
         else {
           PsiElement toImport;
           if (importStatement instanceof PyFromImportStatement) {
-            toImport = ((PyFromImportStatement) importStatement).resolveImportSource();
+            toImport = ((PyFromImportStatement)importStatement).resolveImportSource();
           }
           else {
             toImport = importStatement.getImportElements()[0].resolve();
@@ -140,12 +143,21 @@ public class PyImportOptimizer implements ImportOptimizer {
       addImports(myThirdPartyImports);
       addImports(myProjectImports);
       PsiElement lastElement = myImportBlock.get(myImportBlock.size()-1);
-      myFile.deleteChildRange(myImportBlock.get(0), lastElement);
+      myFile.deleteChildRange(findFirstNonFutureImport(), lastElement);
       for (PyImportStatementBase anImport : myBuiltinImports) {
         anImport.putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, null);
       }
     }
 
+    private PyImportStatementBase findFirstNonFutureImport() {
+      for (PyImportStatementBase importStatement: myImportBlock) {
+        if (!(importStatement instanceof PyFromImportStatement && ((PyFromImportStatement)importStatement).isFromFuture())) {
+          return importStatement;
+        }
+      }
+      return myImportBlock.get(0);
+    }
+
     private static void markGroupBegin(List<PyImportStatementBase> imports) {
       if (imports.size() > 0) {
         imports.get(0).putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, true);
@@ -154,7 +166,7 @@ public class PyImportOptimizer implements ImportOptimizer {
 
     private void addImports(final List<PyImportStatementBase> imports) {
       for (PyImportStatementBase newImport: imports) {
-        myFile.addBefore(newImport, myImportBlock.get(0));
+        myFile.addBefore(newImport, findFirstNonFutureImport());
       }
     }
   }

python/testData/optimizeImports/order.after.py

@@ -1,3 +1,4 @@
+from __future__ import with_statement
 import sys
 import datetime
 

python/testData/optimizeImports/order.py

@@ -1,3 +1,4 @@
+from __future__ import with_statement
 import foo
 import sys
 from bar import *