fix_trace_division: support inline operations

my2to3/fixes/fix_trace_division.py

@@ -3,6 +3,8 @@ from collections import defaultdict
 import lib2to3.fixer_base
 import lib2to3.fixer_util
 import lib2to3.pgen2
+from lib2to3.pygram import python_symbols as syms
+from lib2to3.pytree import Node
 import sys
 
 
@@ -25,34 +27,54 @@ class FixTraceDivision(lib2to3.fixer_base.BaseFix):
 
       import __builtin__
       __builtin__.division_traced = division_traced
-
   """
   def start_tree(self, tree, filename):
     super(FixTraceDivision, self).start_tree(tree, filename)
     self.ids = defaultdict(int)
 
   def match(self, node):
-    return (
-        len(node.children) == 3 and
-        node.children[1].type == lib2to3.pgen2.token.SLASH)
+    if node.type == syms.term:
+      # List the position of '/' nodes
+      return [i for i, child in enumerate(node.children)
+              if child.type == lib2to3.pgen2.token.SLASH]
 
   def transform(self, node, results):
+    previous_node = None
+
+    for operator_pos in range(1, len(node.children), 2):
+
+      children = [leaf 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]
+
+      if operator_pos in results:
+        # It's a division operation
         lineno = node.get_lineno()
-    dividend = node.children[0].clone()
+        dividend = children[0]
         if node.prefix == dividend.prefix:
           dividend.prefix = ''
         comma = lib2to3.fixer_util.Comma()
-    comma.prefix = prefix = node.children[1].prefix
+        comma.prefix = children[1].prefix
 
-    node.replace(
-        lib2to3.fixer_util.Call(
+        new_node = lib2to3.fixer_util.Call(
           lib2to3.fixer_util.Name("division_traced"),
           args=(
             lib2to3.fixer_util.Number(self.ids[lineno]),
             lib2to3.fixer_util.Comma(),
             dividend,
             comma,
-                node.children[2].clone(),
+            children[2].clone(),
           ),
-            prefix=node.prefix))
+          prefix=node.prefix)
+
         self.ids[lineno] += 1
+
+      else:
+        # It's not a division operation
+        new_node = Node(syms.term, children)
+
+      previous_node = new_node
+
+    node.replace(previous_node)

my2to3/tests/testFixTraceDivision.py

@@ -26,6 +26,21 @@ class testFixTraceDivision(FixerTestCase):
     a = """division_traced(1,x , (division_traced(0,y , z)))"""
     self.check(b, a)
 
+  def test_inline_division_1(self):
+    b = """1 / 2 / 3"""
+    a = """division_traced(1,division_traced(0,1 , 2) , 3)"""
+    self.check(b, a)
+
+  def test_inline_division_2(self):
+    b = """1 / 2 * 3"""
+    a = """division_traced(0,1 , 2) * 3"""
+    self.check(b, a)
+
+  def test_inline_division_3(self):
+    b = """1 * 2 / 3"""
+    a = """division_traced(0,1 * 2 , 3)"""
+    self.check(b, a)
+
   def test_multiline_division_1(self):
     b = """x \
   / y"""