From 421644134010d413a623bd3f30c59f0cc1ba76dc Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Thu, 31 Oct 2013 17:49:34 +0100
Subject: [PATCH] import sorter skips from __future__ imports (PY-10022)

---
 .../codeInsight/imports/PyImportOptimizer.java | 18 +++++++++++++++---
 python/testData/optimizeImports/order.after.py |  1 +
 python/testData/optimizeImports/order.py       |  1 +
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java
index fee8c8518ce7..6d5e3c850de2 100644
--- a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java
+++ b/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());
       }
     }
   }
diff --git a/python/testData/optimizeImports/order.after.py b/python/testData/optimizeImports/order.after.py
index e89fc2619475..e3719e04dcc0 100644
--- a/python/testData/optimizeImports/order.after.py
+++ b/python/testData/optimizeImports/order.after.py
@@ -1,3 +1,4 @@
+from __future__ import with_statement
 import sys
 import datetime
 
diff --git a/python/testData/optimizeImports/order.py b/python/testData/optimizeImports/order.py
index 4232fbd83ed6..ad3147aef522 100644
--- a/python/testData/optimizeImports/order.py
+++ b/python/testData/optimizeImports/order.py
@@ -1,3 +1,4 @@
+from __future__ import with_statement
 import foo
 import sys
 from bar import *