Collapse else: if: ... into elif:

This commit is contained in:
Mark Dickinson 2010-06-30 08:32:11 +00:00
parent 719e4e3ba5
commit 8d6d760422
2 changed files with 32 additions and 1 deletions

View File

@ -64,6 +64,24 @@ class_decorator = """\
class Foo: pass
"""
elif1 = """\
if cond1:
suite1
elif cond2:
suite2
else:
suite3
"""
elif2 = """\
if cond1:
suite1
elif cond2:
suite2
"""
class ASTTestCase(unittest.TestCase):
def assertASTEqual(self, ast1, ast2):
self.assertEqual(ast.dump(ast1), ast.dump(ast2))
@ -159,6 +177,10 @@ class UnparseTestCase(ASTTestCase):
def test_class_definition(self):
self.check_roundtrip("class A(metaclass=type, *[], **{}): pass")
def test_elifs(self):
self.check_roundtrip(elif1)
self.check_roundtrip(elif2)
class DirectoryTestCase(ASTTestCase):
"""Test roundtrip behaviour on all files in Lib and Lib/test."""

View File

@ -256,9 +256,18 @@ class Unparser:
self.fill("if ")
self.dispatch(t.test)
self.enter()
# XXX elif?
self.dispatch(t.body)
self.leave()
# collapse nested ifs into equivalent elifs.
while (t.orelse and len(t.orelse) == 1 and
isinstance(t.orelse[0], ast.If)):
t = t.orelse[0]
self.fill("elif ")
self.dispatch(t.test)
self.enter()
self.dispatch(t.body)
self.leave()
# final else
if t.orelse:
self.fill("else")
self.enter()