Import statements should always be placed after any comments or whitespace.

Addresses issue #18

Imports are already placed after copyright/comments if there's a
docstring or shebang in the file. This makes import handling consistent in more cases.

Author: jeremyh

libmodernize/__init__.py

@@ -61,10 +61,9 @@ def add_future(node, symbol):
     import_ = fixer_util.FromImport('__future__',
                                     [Leaf(token.NAME, symbol, prefix=" ")])
 
-    # If we're inserting as the first element, ensure any shebang prefix is maintained.
-    if idx == 0 and node.prefix.startswith('#!'):
-        import_.prefix = node.prefix
-        node.prefix = ''
+    # Place after any comments or whitespace. (copyright, shebang etc.)
+    import_.prefix = node.prefix
+    node.prefix = ''
 
     children = [import_, fixer_util.Newline()]
     root.insert_child(idx, Node(syms.simple_stmt, children))

tests/test_fix_import.py

@@ -30,14 +30,14 @@
 """)
 
 DOCSTRING = ("""\
-\"\"\"
+\"""
 Docstring
-\"\"\"
+\"""
 import foo
 """, """\
-\"\"\"
+\"""
 Docstring
-\"\"\"
+\"""
 from __future__ import absolute_import
 import foo
 """)
@@ -53,15 +53,15 @@
 
 DOCSTING_AND_SHEBANG = ("""\
 #!/usr/bin/env python
-\"\"\"
+\"""
 Docstring
-\"\"\"
+\"""
 import foo
 """, """\
 #!/usr/bin/env python
-\"\"\"
+\"""
 Docstring
-\"\"\"
+\"""
 from __future__ import absolute_import
 import foo
 """)
@@ -85,6 +85,27 @@
 import foo
 """)
 
+
+COPYRIGHT_AND_DOCSTRING = ("""\
+#
+# Copyright notice
+#
+
+\"""Docstring\"""
+
+import foo
+""", """\
+#
+# Copyright notice
+#
+
+\"""Docstring\"""
+
+from __future__ import absolute_import
+import foo
+""")
+
+
 def test_no_imports():
     check_on_input(*NO_IMPORTS)
 
@@ -108,3 +129,6 @@ def test_import_with_docstring_and_shebang():
 
 def test_import_with_copyright_and_shebang():
     check_on_input(*COPYRIGHT_AND_SHEBANG)
+
+def test_import_with_copyright_and_docstring():
+    check_on_input(*COPYRIGHT_AND_DOCSTRING)