Python runtime bytecode patch example

shame.py

import functools
import struct
import types
import dis


def _patch_code(code: types.CodeType):
    # Call function with 0 arguments
    function_call_bytearray = bytearray(dis.opmap['CALL_FUNCTION'].to_bytes(1, byteorder='little') + b'\x00')

    code_bytarray = bytearray(code.co_code)
    names_list = list(code.co_names)

    # I believe each python bytecode command has 2 bytes
    # Range is not working successfully with 3 arguments (what)
    for i in range(0, len(code.co_code) // 2):
        i *= 2
        opcode, oparg = struct.unpack_from('BB', code_bytarray, i)
        if dis.opname[opcode] != 'LOAD_GLOBAL':
            continue

        global_name = names_list[oparg]
        if global_name != 'shame':
            continue

        names_list[oparg] = 'print_shame'
        # Insert CALL_FUNCTION after LOAD_GLOBAL 'shame'
        code_bytarray[i + 2:i + 2] = function_call_bytearray

    return _make_code(code, bytes(code_bytarray), tuple(names_list))


def _make_code(code, codestring, names):
    # Successfully snatched from
    # https://github.com/snoack/python-goto/blob/master/goto.py#L45
    args = [
        code.co_argcount, code.co_nlocals, code.co_stacksize,
        code.co_flags, codestring, code.co_consts,
        names, code.co_varnames, code.co_filename,
        code.co_name, code.co_firstlineno, code.co_lnotab,
        code.co_freevars, code.co_cellvars
    ]

    try:
        args.insert(1, code.co_kwonlyargcount)  # PY3
    except AttributeError:
        pass

    return types.CodeType(*args)


def with_shame(func_or_code):
    if isinstance(func_or_code, types.CodeType):
        return _patch_code(func_or_code)

    return functools.update_wrapper(
        types.FunctionType(
            _patch_code(func_or_code.__code__),
            func_or_code.__globals__,
            func_or_code.__name__,
            func_or_code.__defaults__,
            func_or_code.__closure__,
        ),
        func_or_code
    )


def print_shame():
    print('it works lmao bottom text')


@with_shame
def test():
    """
    Will be compiled and executed successfully
    """
    print('> Ok')
    shame
    return 1


if __name__ == '__main__':
    test()