Index: pytree.py =================================================================== --- pytree.py (revision 57363) +++ pytree.py (working copy) @@ -148,6 +148,25 @@ self.parent = None return i + def get_previous_sibling(self): + """Return the node immediately preceding the invocant in their + parent's children list. If the invocant does not have a previous + sibling, return None.""" + if self.parent is None: + return None + + # Can't use index(); we need to test by identity + for i, sibling in enumerate(self.parent.children): + if sibling is self: + # The first child has no previous sibling. Special-case this + # because i-1 would return -1, which is a valid index. + if i == 0: + return None + try: + return self.parent.children[i-1] + except IndexError: + return None + def get_next_sibling(self): """Return the node immediately following the invocant in their parent's children list. If the invocant does not have a next @@ -163,6 +182,16 @@ except IndexError: return None + def get_previous_in_tree(self): + """Return the node immediately preceding the invocant, following + reverse pre-order traversal.""" + sib = self.get_previous_sibling() + if sib is not None: + while sib.children: + sib = sib.children[-1] + return sib + return self.parent + def get_suffix(self): """Return the string immediately following the invocant node. This is effectively equivalent to node.get_next_sibling().get_prefix()""" Index: tests/test_util.py =================================================================== --- tests/test_util.py (revision 57363) +++ tests/test_util.py (working copy) @@ -60,6 +60,116 @@ self.failIf(self.is_list("[]+[]")) +class Test_insert_comment(support.TestCase): + def test_simple(self): + """ + Input Output + ----- ------ + + foo # test + foo + """ + node = parse('foo') + util.insert_comment(node.children[0].children[0], 'test') + self.assertEqual(str(node), "# test\nfoo\n\n") + + def test_simple2(self): + """ + Input Output + ----- ------ + + foo # test + foo + """ + node = parse('foo') + util.insert_comment(node, 'test') + self.assertEqual(str(node), "# test\nfoo\n\n") + + def test_indent1(self): + """ + Input Output + ----- ------ + + if foo: if foo: + bar() # test + bar() + """ + node = parse('if foo:\n bar()') + expected = 'if foo:\n # test\n bar()\n\n' + leaf = node.children[0].children[3].children[2] + self.assertEqual(str(leaf), 'bar()\n') + util.insert_comment(leaf, 'test') + self.assertEqual(str(node), expected) + + def test_indent2(self): + """ + Input Output + ----- ------ + + if foo: if foo: + if bar: if bar: + baz() # test + baz() + """ + node = parse('if foo:\n if bar:\n baz()') + expected = 'if foo:\n if bar:\n # test\n baz()\n\n' + leaf = node.children[0].children[3].children[2].children[3].children[2] + self.assertEqual(str(leaf), 'baz()\n') + util.insert_comment(leaf, 'test') + self.assertEqual(str(node), expected) + + def test_indent3(self): + """ + Input Output + ----- ------ + + if foo: if foo: + bar() bar() + else: else: + baz() # test + baz() + """ + node = parse('if foo:\n bar()\nelse:\n baz()') + expected = 'if foo:\n bar()\nelse:\n # test\n baz()\n\n' + leaf = node.children[0].children[6].children[2] + self.assertEqual(str(leaf), 'baz()\n') + util.insert_comment(leaf, 'test') + self.assertEqual(str(node), expected) + + def test_indent4(self): + """ + Input Output + ----- ------ + + if foo: if foo: + bar() bar() + else: # test + baz() else: + baz() + """ + node = parse('if foo:\n bar()\nelse:\n baz()') + expected = 'if foo:\n bar()\n# test\nelse:\n baz()\n\n' + leaf = node.children[0].children[4] + self.assertEqual(str(leaf), 'else') + util.insert_comment(leaf, 'test') + self.assertEqual(str(node), expected) + + def test_indent5(self): + """ + Input Output + ----- ------ + + # legacycomment # legacycomment + foo # test + foo + """ + node = parse('# legacycomment\nfoo') + expected = '# legacycomment\n# test\nfoo\n\n' + leaf = node.children[0].children[0] + self.assertEqual(str(leaf), '# legacycomment\nfoo') + util.insert_comment(leaf, 'test') + self.assertEqual(str(node), expected) + class Test_Attr(MacroTestCase): def test(self): from fixes.util import Attr, Name Index: tests/test_pytree.py =================================================================== --- tests/test_pytree.py (revision 57363) +++ tests/test_pytree.py (working copy) @@ -301,6 +301,15 @@ # I don't care what it raises, so long as it's an exception self.assertRaises(Exception, n1.append_child, list) + def testNodePreviousSibling(self): + n1 = pytree.Node(1000, []) + n2 = pytree.Node(1000, []) + p1 = pytree.Node(1000, [n1, n2]) + + self.failUnless(n2.get_previous_sibling() is n1) + self.assertEqual(n1.get_previous_sibling(), None) + self.assertEqual(p1.get_previous_sibling(), None) + def testNodeNextSibling(self): n1 = pytree.Node(1000, []) n2 = pytree.Node(1000, []) @@ -310,6 +319,15 @@ self.assertEqual(n2.get_next_sibling(), None) self.assertEqual(p1.get_next_sibling(), None) + def testLeafPreviousSibling(self): + l1 = pytree.Leaf(100, "a") + l2 = pytree.Leaf(100, "b") + p1 = pytree.Node(1000, [l1, l2]) + + self.failUnless(l2.get_previous_sibling() is l1) + self.assertEqual(l1.get_previous_sibling(), None) + self.assertEqual(p1.get_previous_sibling(), None) + def testLeafNextSibling(self): l1 = pytree.Leaf(100, "a") l2 = pytree.Leaf(100, "b") @@ -319,7 +337,46 @@ self.assertEqual(l2.get_next_sibling(), None) self.assertEqual(p1.get_next_sibling(), None) + def testGetPreviousInTree(self): + # Create the following tree: + # + # P1 + # / \ + # N1 N2 + # / \ / \ + # N3 N4 N5 N6 + # / / | / \ + # N7 N8 N9 N10 N11 + # | + # N12 + n7 = pytree.Node(1000, []) + n8 = pytree.Node(1000, []) + n9 = pytree.Node(1000, []) + n5 = pytree.Node(1000, []) + n12 = pytree.Node(1000, []) + n11 = pytree.Node(1000, []) + n3 = pytree.Node(1000, [n7]) + n4 = pytree.Node(1000, [n8, n9]) + n10 = pytree.Node(1000, [n12]) + n6 = pytree.Node(1000, [n10, n11]) + n1 = pytree.Node(1000, [n3, n4]) + n2 = pytree.Node(1000, [n5, n6]) + p1 = pytree.Node(1000, [n1, n2]) + self.assertEqual(p1.get_previous_in_tree(), None) + self.failUnless(n1.get_previous_in_tree() is p1) + self.failUnless(n2.get_previous_in_tree() is n9) + self.failUnless(n3.get_previous_in_tree() is n1) + self.failUnless(n4.get_previous_in_tree() is n7) + self.failUnless(n5.get_previous_in_tree() is n2) + self.failUnless(n6.get_previous_in_tree() is n5) + self.failUnless(n7.get_previous_in_tree() is n3) + self.failUnless(n8.get_previous_in_tree() is n4) + self.failUnless(n9.get_previous_in_tree() is n8) + self.failUnless(n10.get_previous_in_tree() is n6) + self.failUnless(n11.get_previous_in_tree() is n12) + self.failUnless(n12.get_previous_in_tree() is n10) + class TestPatterns(support.TestCase): """Unit tests for tree matching patterns.""" Index: fixes/util.py =================================================================== --- fixes/util.py (revision 57363) +++ fixes/util.py (working copy) @@ -103,6 +103,39 @@ inner, Leaf(token.RBRACE, "]")]) +def insert_comment(node, comment_text): + """Inserts a Python comment on the line above the given Node or Leaf, + taking indendation into account.""" + + # Calculate the indentation level of the given node. Do this by traversing + # backwards (previous sibling(s), then parent) until either a token.NEWLINE + # or token.INDENT is encountered. When this loop is done, `indent` will + # hold a string representing the exact indentation, and `node` will be the + # node immediately after the newline -- the one before which the comment + # should be added. + + indent = '' + while node.parent is not None: + prev = node.get_previous_in_tree() + if prev.type == token.INDENT: + indent = prev.value + node = prev + break + elif prev.type == token.NEWLINE: + break + node = prev + + # Drill down into the Node until we have a Leaf. + while isinstance(node, Node): + node = node.children[0] + + # Comments are represented in the AST using the 'prefix' attribute + # of a Leaf, so we add comments by changing the 'prefix' attribute. + # Note that we're including the previous node.prefix in here, in case + # this Leaf already had a comment on it. In that case, our new comment + # will be added *after* it. + node.set_prefix('%s%s# %s\n' % (indent, node.prefix, comment_text)) + ########################################################### ### Determine whether a node represents a given literal ###########################################################