Import statements should be placed after the shebang.

Addresses issue #18

Imports are also placed after any comments (such as copyright notices), to be consistent with the existing import placement in files with module docstrings.

Author: jeremyh

libmodernize/__init__.py

@@ -60,6 +60,12 @@ 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 = ''
+
     children = [import_, fixer_util.Newline()]
     root.insert_child(idx, Node(syms.simple_stmt, children))
 

tests/test_fix_import.py

@@ -29,6 +29,61 @@
 import foo
 """)
 
+DOCSTRING = ("""\
+\"\"\"
+Docstring
+\"\"\"
+import foo
+""", """\
+\"\"\"
+Docstring
+\"\"\"
+from __future__ import absolute_import
+import foo
+""")
+
+SHEBANG = ("""\
+#!/usr/bin/env python
+import foo
+""", """\
+#!/usr/bin/env python
+from __future__ import absolute_import
+import foo
+""")
+
+DOCSTING_AND_SHEBANG = ("""\
+#!/usr/bin/env python
+\"\"\"
+Docstring
+\"\"\"
+import foo
+""", """\
+#!/usr/bin/env python
+\"\"\"
+Docstring
+\"\"\"
+from __future__ import absolute_import
+import foo
+""")
+
+COPYRIGHT_AND_SHEBANG = ("""\
+#!/usr/bin/env python
+
+#
+# Copyright notice
+#
+
+import foo
+""", """\
+#!/usr/bin/env python
+
+#
+# Copyright notice
+#
+
+from __future__ import absolute_import
+import foo
+""")
 
 def test_no_imports():
     check_on_input(*NO_IMPORTS)
@@ -41,3 +96,15 @@ def test_only_normal_imports():
 
 def test_normal_and_future_imports():
     check_on_input(*NORMAL_AND_FUTURE_IMPORTS)
+
+def test_import_with_docstring():
+    check_on_input(*DOCSTRING)
+
+def test_import_with_shebang():
+    check_on_input(*SHEBANG)
+
+def test_import_with_docstring_and_shebang():
+    check_on_input(*DOCSTING_AND_SHEBANG)
+
+def test_import_with_copyright_and_shebang():
+    check_on_input(*COPYRIGHT_AND_SHEBANG)