# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
# For details: https://bitbucket.org/ned/coveragepy/src/default/NOTICE.txt

""" Wicked hack to get .pyc files to do bytecode tracing instead of
    line tracing.
"""

import marshal, new, opcode, sys, types

from lnotab import lnotab_numbers, lnotab_string

class PycFile:
    def read(self, f):
        if isinstance(f, basestring):
            f = open(f, "rb")
        self.magic = f.read(4)
        self.modtime = f.read(4)
        self.code = marshal.load(f)

    def write(self, f):
        if isinstance(f, basestring):
            f = open(f, "wb")
        f.write(self.magic)
        f.write(self.modtime)
        marshal.dump(self.code, f)

    def hack_line_numbers(self):
        self.code = hack_line_numbers(self.code)

def hack_line_numbers(code):
    """ Replace a code object's line number information to claim that every
        byte of the bytecode is a new source line.  Returns a new code
        object.  Also recurses to hack the line numbers in nested code objects.
    """

    # Create a new lnotab table.  Each opcode is claimed to be at
    # 1000*lineno + (opcode number within line), so for example, the opcodes on
    # source line 12 will be given new line numbers 12000, 12001, 12002, etc.
    old_num = list(lnotab_numbers(code.co_lnotab, code.co_firstlineno))
    n_bytes = len(code.co_code)
    new_num = []
    line = 0
    opnum_in_line = 0
    i_byte = 0
    while i_byte < n_bytes:
        if old_num and i_byte == old_num[0][0]:
            line = old_num.pop(0)[1]
            opnum_in_line = 0
        new_num.append((i_byte, 100000000 + 1000*line + opnum_in_line))
        if ord(code.co_code[i_byte]) >= opcode.HAVE_ARGUMENT:
            i_byte += 3
        else:
            i_byte += 1
        opnum_in_line += 1

    # new_num is a list of pairs, (byteoff, lineoff).  Turn it into an lnotab.
    new_firstlineno = new_num[0][1]-1
    new_lnotab = lnotab_string(new_num, new_firstlineno)

    # Recurse into code constants in this code object.
    new_consts = []
    for const in code.co_consts:
        if type(const) == types.CodeType:
            new_consts.append(hack_line_numbers(const))
        else:
            new_consts.append(const)

    # Create a new code object, just like the old one, except with new
    # line numbers.
    new_code = new.code(
        code.co_argcount, code.co_nlocals, code.co_stacksize, code.co_flags,
        code.co_code, tuple(new_consts), code.co_names, code.co_varnames,
        code.co_filename, code.co_name, new_firstlineno, new_lnotab
        )

    return new_code

def hack_file(f):
    pyc = PycFile()
    pyc.read(f)
    pyc.hack_line_numbers()
    pyc.write(f)

if __name__ == '__main__':
    hack_file(sys.argv[1])