from test_utils import *


def test_constant():

    @Commentor("<return>")
    def target():
        x = 2
        print(x == 2)

    asserteq_or_print(target(), '''
    def target():
        x = 2
        print(x == 2)
        """
        <callable> : print
        True : x == 2
        None : print(x == 2)
        """
''')


def test_if():

    @Commentor("<return>")
    def target():
        x = 2
        if x > 3:
            x = 2 * x
            y = 1
        elif x > 2:
            x = 4 * x
            y = 2
        elif x > 3:
            x = 4 * x
            y = 3
        else:
            x = 8 * x
            y = 5
    
    asserteq_or_print(target(), '''
    def target():
        x = 2
        if x > 3:  # False
            x = 2 * x  # skipped
            y = 1  # skipped
        elif x > 2:  # False
            x = 4 * x  # skipped
            y = 2  # skipped
        elif x > 3:  # False
            x = 4 * x  # skipped
            y = 3  # skipped
        else:    # True
            x = 8 * x
            """
            2 : x
            16 : 8 * x
            ----------
            16 : x
            """
            y = 5
''')


def test_for():
    
    with closing(StringIO()) as f:
        @Commentor(f, _exit=False, check=False)
        def target():
            odds = []
            # return only odd numbers - 1,3,5,7,9
            for x in range(10):
                # Check if x is even
                if x % 2 == 0:
                    continue
                odds.append(x)
            return odds
        
        assert target() == [1, 3, 5, 7, 9]
        asserteq_or_print(f.getvalue(), '''
        def target():
            odds = []
            for x in range(10):
                
                ###### !new iteration! ######
                """
                0 : __REG__for_loop_iter_once
                ----------
                0 : x
                """
                # if x % 2 == 0:  # True
                #     continue  # True
                # odds.append(x)  # skipped
                
                ###### !new iteration! ######
                """
                1 : __REG__for_loop_iter_once
                ----------
                1 : x
                """
                # if x % 2 == 0:  # False
                #     continue  # skipped
                # odds.append(x)
                """
                [] : odds
                <callable> : odds.append
                1 : x
                None : odds.append(x)
                """
                
                ###### !new iteration! ######
                """
                2 : __REG__for_loop_iter_once
                ----------
                2 : x
                """
                # if x % 2 == 0:  # True
                #     continue  # True
                # odds.append(x)  # skipped
                
                ###### !new iteration! ######
                """
                3 : __REG__for_loop_iter_once
                ----------
                3 : x
                """
                # if x % 2 == 0:  # False
                #     continue  # skipped
                # odds.append(x)
                """
                [1] : odds
                <callable> : odds.append
                3 : x
                None : odds.append(x)
                """
                
                ###### !new iteration! ######
                """
                4 : __REG__for_loop_iter_once
                ----------
                4 : x
                """
                # if x % 2 == 0:  # True
                #     continue  # True
                # odds.append(x)  # skipped
                
                ###### !new iteration! ######
                """
                5 : __REG__for_loop_iter_once
                ----------
                5 : x
                """
                # if x % 2 == 0:  # False
                #     continue  # skipped
                # odds.append(x)
                """
                [1, 3] : odds
                <callable> : odds.append
                5 : x
                None : odds.append(x)
                """
                
                ###### !new iteration! ######
                """
                6 : __REG__for_loop_iter_once
                ----------
                6 : x
                """
                # if x % 2 == 0:  # True
                #     continue  # True
                # odds.append(x)  # skipped
                
                ###### !new iteration! ######
                """
                7 : __REG__for_loop_iter_once
                ----------
                7 : x
                """
                # if x % 2 == 0:  # False
                #     continue  # skipped
                # odds.append(x)
                """
                [1, 3, 5] : odds
                <callable> : odds.append
                7 : x
                None : odds.append(x)
                """
                
                ###### !new iteration! ######
                """
                8 : __REG__for_loop_iter_once
                ----------
                8 : x
                """
                # if x % 2 == 0:  # True
                #     continue  # True
                # odds.append(x)  # skipped
                
                ###### !new iteration! ######
                """
                9 : __REG__for_loop_iter_once
                ----------
                9 : x
                """
                if x % 2 == 0:  # False
                    continue  # skipped
                odds.append(x)
                """
                [1, 3, 5, 7] : odds
                <callable> : odds.append
                9 : x
                None : odds.append(x)
                """
            return odds
            """
            [1, 3, 5, 7, 9] : odds
            """''')