adhoc dtrace in python

pytrace.py

#!/usr/bin/env python

"""
dtrace -n return'{ustack()}'
for python
"""

from byteplay import *

def seekret(code):
    def gen():
        lineno = 0

        for i, (op, arg) in enumerate(code):
            if op == SetLineno:
                lineno = i + 1
            elif op == RETURN_VALUE:
                yield lineno

    return list(gen())

def codesplit(code, patchpoints, last=0):
    if not patchpoints:
        return [code[last:]]

    point = patchpoints.pop(0)
    chunk = code[last:point]

    return ([chunk] if chunk else []) + codesplit(code, patchpoints, last=point)

TRACEBACK_PATCH = [
    (LOAD_CONST, -1),
    (LOAD_CONST, None),
    (IMPORT_NAME, 'traceback'),
    (STORE_FAST, 'traceback'),
    (LOAD_FAST, 'traceback'),
    (LOAD_ATTR, 'print_stack'),
    (CALL_FUNCTION, 0),
    (POP_TOP, None),
]

def inject(patch, chunks):
    def gen():
        yield chunks[0]
        for code in chunks[1:]:
            yield patch
            yield code
    return reduce(list.__add__, gen(), [])

def patch_pystackret(fun):
    assert hasattr(fun, 'func_code')
    if getattr(fun, '_pystackret_patched', False):
        return

    # disassemble
    codeobj = Code.from_code(fun.func_code)
    code = codeobj.code
    # patch
    codeobj.code = inject(TRACEBACK_PATCH, codesplit(code, seekret(code)))

    #from pprint import pprint
    #print pprint(codeobj.code)

    # assemble
    fun.func_code = codeobj.to_code()

    fun._pystackret_patched = True


if __name__ == '__main__':
    def acme(arg):
        if arg:
            return 1
        else:
            return 0

    patch_pystackret(acme)
    print '::> calling acme(True)'
    acme(True)
    print '::> calling acme(False)'
    acme(False)

    """
::> calling acme(True)
  File "./pytrace.py", line 79, in <module>
    acme(True)
  File "./pytrace.py", line 73, in acme
    return 1
::> calling acme(False)
  File "./pytrace.py", line 81, in <module>
    acme(False)
  File "./pytrace.py", line 75, in acme
    return 0
"""