def test_if_no_ifs(self):
scope = codegen.Module()
if_statement = codegen.If(scope)
if_statement.else_block.add_return(codegen.Number(3))
if_statement = if_statement.finalize()
self.assertCodeEqual(as_source_code(if_statement), """
return 3
""")
def test_if_one_if(self):
scope = codegen.Module()
if_statement = codegen.If(scope)
first_block = if_statement.add_if(codegen.Number(1))
first_block.add_return(codegen.Number(2))
self.assertCodeEqual(as_source_code(if_statement), """
if 1:
return 2
""")
def test_if_two_ifs(self):
scope = codegen.Module()
if_statement = codegen.If(scope)
first_block = if_statement.add_if(codegen.Number(1))
first_block.add_return(codegen.Number(2))
second_block = if_statement.add_if(codegen.Number(3))
second_block.add_return(codegen.Number(4))
self.assertCodeEqual(as_source_code(if_statement), """
if 1:
return 2
elif 3:
return 4
""")
def test_if_empty(self):
scope = codegen.Module()
if_statement = codegen.If(scope)
if_statement = if_statement.finalize()
self.assertCodeEqual(as_source_code(if_statement), "")