"""bytecode_helper - support tools for testing correct bytecode generation"""

import unittest
import dis
import io

_UNSPECIFIED = object()

class BytecodeTestCase(unittest.TestCase):
    """Custom assertion methods for inspecting bytecode."""

    def get_disassembly_as_string(self, co):
        s = io.StringIO()
        dis.dis(co, file=s)
        return s.getvalue()

    def assertInstructionMatches(self, instr, expected, *, line_offset=0):
        # Deliberately test opname first, since that gives a more
        # meaningful error message than testing opcode
        self.assertEqual(instr.opname, expected.opname)
        self.assertEqual(instr.opcode, expected.opcode)
        self.assertEqual(instr.arg, expected.arg)
        self.assertEqual(instr.argval, expected.argval)
        self.assertEqual(instr.argrepr, expected.argrepr)
        self.assertEqual(instr.offset, expected.offset)
        if expected.starts_line is None:
            self.assertIsNone(instr.starts_line)
        else:
            self.assertEqual(instr.starts_line,
                                expected.starts_line + line_offset)
        self.assertEqual(instr.is_jump_target, expected.is_jump_target)


    def assertBytecodeExactlyMatches(self, x, expected, *, line_offset=0):
        """Throws AssertionError if any discrepancy is found in bytecode

        *x* is the object to be introspected
        *expected* is a list of dis.Instruction objects

        Set *line_offset* as appropriate to adjust for the location of the
        object to be disassembled within the test file. If the expected list
        assumes the first line is line 1, then an appropriate offset would be
        ``1 - f.__code__.co_firstlineno``.
        """
        actual = dis.get_instructions(x, line_offset=line_offset)
        self.assertEqual(list(actual), expected)

    def assertInBytecode(self, x, opname, argval=_UNSPECIFIED):
        """Returns instr if op is found, otherwise throws AssertionError"""
        for instr in dis.get_instructions(x):
            if instr.opname == opname:
                if argval is _UNSPECIFIED or instr.argval == argval:
                    return instr
        disassembly = self.get_disassembly_as_string(x)
        if argval is _UNSPECIFIED:
            msg = '%s not found in bytecode:\n%s' % (opname, disassembly)
        else:
            msg = '(%s,%r) not found in bytecode:\n%s'
            msg = msg % (opname, argval, disassembly)
        self.fail(msg)

    def assertNotInBytecode(self, x, opname, argval=_UNSPECIFIED):
        """Throws AssertionError if op is found"""
        for instr in dis.get_instructions(x):
            if instr.opname == opname:
                disassembly = self.get_disassembly_as_string(co)
                if opargval is _UNSPECIFIED:
                    msg = '%s occurs in bytecode:\n%s' % (opname, disassembly)
                elif instr.argval == argval:
                    msg = '(%s,%r) occurs in bytecode:\n%s'
                    msg = msg % (opname, argval, disassembly)
                self.fail(msg)