New fixer: round_trace

Also, some code refactoring.

my2to3/fixes/__init__.py

+from collections import defaultdict
+from lib2to3.fixer_base import BaseFix as lib2to3_BaseFix
+import lib2to3.fixer_util
+
+
+class BaseFix(lib2to3_BaseFix):
+  def start_tree(self, tree, filename):
+    super(BaseFix, self).start_tree(tree, filename)
+
+    self.ids = defaultdict(int)
+
+  def traced_call(self, name, insert_function, node, children):
+    # Important: every node inside "children" should be cloned before this
+    # function is called.
+
+    first_child = children[0]
+    if node.prefix == first_child.prefix:
+      first_child.prefix = ''
+
+    filename = self.filename
+    lineno = node.get_lineno()
+    id_ = self.ids[lineno]
+
+    new_node = lib2to3.fixer_util.Call(
+      lib2to3.fixer_util.Name(name),
+      args=[
+        lib2to3.fixer_util.String('"%s"' % filename),
+        lib2to3.fixer_util.Comma(),
+        lib2to3.fixer_util.Number(lineno),
+        lib2to3.fixer_util.Comma(),
+        lib2to3.fixer_util.Number(id_),
+        lib2to3.fixer_util.Comma(),
+      ] + children,
+      prefix=node.prefix)
+
+    self.ids[lineno] += 1
+
+    insert_function(filename, lineno, id_)
+
+    return new_node

my2to3/fixes/fix_division_support.py

-from lib2to3.fixer_base import BaseFix
 import lib2to3.pgen2
 from lib2to3.pygram import python_symbols as syms
 from lib2to3.pytree import Leaf, Node
 import os
 
+from . import BaseFix
 from my2to3.trace import create_table, get_data
 from my2to3.util import add_future, data2types
 

my2to3/fixes/fix_division_trace.py

 # https://lab.nexedi.com/nexedi/erp5/snippets/475
-from collections import defaultdict
-from lib2to3.fixer_base import BaseFix
 import lib2to3.fixer_util
 import lib2to3.pgen2
 from lib2to3.pygram import python_symbols as syms
 from lib2to3.pytree import Node
 
+from . import BaseFix
 from my2to3.trace import create_table, register_tracing_function
 from my2to3.util import parse_type
 
 
-# id is used to differentiate the divisions of the same line.
 insert_trace = create_table("division_trace", "filename", "lineno", "id", "dividend_type", "divisor_type")
 insert_modified = create_table("division_modified", "filename", "lineno", "id")
 
@@ -30,13 +28,10 @@ def division_traced(filename, lineno, id_, dividend, divisor):
 class FixDivisionTrace(BaseFix):
   """Rewrites a / b into division_traced(id, a, b)
   """
-  def start_tree(self, tree, filename):
-    super(FixDivisionTrace, self).start_tree(tree, filename)
-    self.ids = defaultdict(int)
 
   def match(self, node):
     if node.type == syms.term:
-      # List the position of '/' nodes
+      # List the position of '/' sign nodes
       return [i for i, child in enumerate(node.children)
               if child.type == lib2to3.pgen2.token.SLASH]
 
@@ -54,45 +49,24 @@ class FixDivisionTrace(BaseFix):
 
     for operator_pos in range(1, len(node.children), 2):
 
-      children = [leaf for leaf in node.children[operator_pos - 1:operator_pos + 2]]
+      # Get the two operators and the operation sign nodes
+      children = [leaf.clone() for leaf in node.children[operator_pos - 1:operator_pos + 2]]
       if previous_node:
         # The first child is actually the previous operation
-        children[0] = previous_node
-      children = [leaf.clone() for leaf in children]
+        children[0] = previous_node.clone()
 
+      # Check if the current operation is a division or not
       if operator_pos in results:
         # It's a division operation
-        lineno = node.get_lineno()
-        dividend = children[0]
-        if node.prefix == dividend.prefix:
-          dividend.prefix = ''
+
+        # Replace the '/' sign by a comma
         comma = lib2to3.fixer_util.Comma()
         comma.prefix = children[1].prefix
+        children[1] = comma
 
-        id_ = self.ids[lineno]
-        insert_modified(self.filename, lineno, id_)
-
-        new_node = lib2to3.fixer_util.Call(
-          lib2to3.fixer_util.Name("division_traced"),
-          args=(
-            lib2to3.fixer_util.String('"%s"' % self.filename),
-            lib2to3.fixer_util.Comma(),
-            lib2to3.fixer_util.Number(lineno),
-            lib2to3.fixer_util.Comma(),
-            lib2to3.fixer_util.Number(id_),
-            lib2to3.fixer_util.Comma(),
-            dividend,
-            comma,
-            children[2].clone(),
-          ),
-          prefix=node.prefix)
-
-        self.ids[lineno] += 1
-
+        previous_node = self.traced_call("division_traced", insert_modified, node, children)
       else:
         # It's not a division operation
-        new_node = Node(syms.term, children)
-
-      previous_node = new_node
+        previous_node = Node(syms.term, children)
 
     node.replace(previous_node)

my2to3/fixes/fix_nested_except_trace.py

-from lib2to3.fixer_base import BaseFix
 from lib2to3.pygram import python_symbols as syms
 
+from . import BaseFix
 from my2to3.trace import create_table
 
 

my2to3/fixes/fix_round_trace.py

+from collections import defaultdict
+import lib2to3.fixer_util
+import lib2to3.pgen2
+from lib2to3.pygram import python_symbols as syms
+from lib2to3.pytree import Node
+
+from . import BaseFix
+from my2to3.trace import create_table, register_tracing_function
+from my2to3.util import parse_type
+
+
+insert_trace = create_table("round_trace", "filename", "lineno", "id", "number", "ndigits")
+insert_modified = create_table("round_modified", "filename", "lineno", "id")
+
+
+@register_tracing_function
+def round_traced(filename, lineno, id_, number, ndigits=0):
+  insert_trace(
+    filename,
+    lineno,
+    id_,
+    number,
+    ndigits
+  )
+  return round(number, ndigits)
+
+
+class FixRoundTrace(BaseFix):
+  """Rewrites round(a[, b]) into round_traced(id, a[, b])
+  """
+  # Inspired by https://github.com/python/cpython/blob/e42b705188271da108de42b55d9344642170aa2b/Lib/lib2to3/fixes/fix_xrange.py#L14
+  PATTERN = """
+            power<'round' trailer< '(' args=any ')' >
+            rest=any* >
+            """
+
+  def transform(self, node, results):
+    args = results['args']
+    if args.type == lib2to3.pgen2.token.NUMBER:
+      # e.g. round(1.23)
+      children = [
+        args.clone()
+      ]
+    else:
+      # e.g. round(1,23, 4)
+      children = [leaf.clone() for leaf in args.children]
+
+    node.replace(self.traced_call("division_traced", insert_modified, node, children))