fix_division: add 'from __future__ import division' when adequate

my2to3/fixes/fix_division.py

@@ -7,6 +7,8 @@ from lib2to3.pytree import Leaf, Node
 import os
 import re
 
+from my2to3.util import add_future
+
 trace_file_match = re.compile(r"^(.*):(.*):(.*) <type '(.*)'> / <type '(.*)'>$").match
 
 
@@ -70,6 +72,7 @@ class FixDivision(lib2to3.fixer_base.BaseFix):
       return
 
     if analyze_data(data):
+      add_future(node, 'division')
       operator = Leaf(lib2to3.pgen2.token.DOUBLESLASH, "//")
     else:
       operator = Leaf(lib2to3.pgen2.token.SLASH, "/")

my2to3/util.py

+from lib2to3 import fixer_util
+from lib2to3.pytree import Leaf, Node
+from lib2to3.pgen2 import token
+from lib2to3.pygram import python_symbols as syms
+
+
+# https://github.com/python-modernize/python-modernize/blob/84d973cb7b8153f9f7f22c3574a59312b2372ccb/libmodernize/__init__.py#L10-49
+def check_future_import(node):
+    """If this is a future import, return set of symbols that are imported,
+    else return None."""
+    # node should be the import statement here
+    if not (node.type == syms.simple_stmt and node.children):
+        return set()
+    node = node.children[0]
+    # now node is the import_from node
+    if not (node.type == syms.import_from and
+            node.children[1].type == token.NAME and
+            node.children[1].value == u'__future__'):
+        return set()
+
+    if node.children[3].type == token.LPAR:
+        # from __future__ import (..
+        node = node.children[4]
+    else:
+        # from __future__ import ...
+        node = node.children[3]
+    # now node is the import_as_name[s]
+
+    # print(python_grammar.number2symbol[node.type])
+    if node.type == syms.import_as_names:
+        result = set()
+        for n in node.children:
+            if n.type == token.NAME:
+                result.add(n.value)
+            elif n.type == syms.import_as_name:
+                n = n.children[0]
+                assert n.type == token.NAME
+                result.add(n.value)
+        return result
+    elif node.type == syms.import_as_name:
+        node = node.children[0]
+        assert node.type == token.NAME
+        return set([node.value])
+    elif node.type == token.NAME:
+        return set([node.value])
+    else:  # pragma: no cover
+        assert 0, "strange import"
+
+
+
+# https://github.com/python-modernize/python-modernize/blob/84d973cb7b8153f9f7f22c3574a59312b2372ccb/libmodernize/__init__.py#L51-L66
+def add_future(node, symbol):
+
+    root = fixer_util.find_root(node)
+
+    for idx, node in enumerate(root.children):
+        if node.type == syms.simple_stmt and \
+           len(node.children) > 0 and node.children[0].type == token.STRING:
+            # skip over docstring
+            continue
+        names = check_future_import(node)
+        if not names:
+            # not a future statement; need to insert before this
+            break
+        if symbol in names:
+            # already imported
+            return
+
+    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))