Merge pull request #123 from jeremyh/master

Import statements should be placed after the shebang.

Author: daira

libmodernize/__init__.py

@@ -60,6 +60,11 @@ def add_future(node, symbol):
 
     import_ = fixer_util.FromImport('__future__',
                                     [Leaf(token.NAME, symbol, 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

@@ -29,6 +29,82 @@
 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
+""")
+
+
+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)
@@ -41,3 +117,18 @@ 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)
+
+def test_import_with_copyright_and_docstring():
+    check_on_input(*COPYRIGHT_AND_DOCSTRING)