测试我们的新操作码

我们现在拥有了注入 DEBUG_OP 的所有工具，让我们来验证下我们的实现是否可用。我们将我们的操作码注入到一个最简单的函数中：

    from new_code import MutableCodeObject  

    def op_target(*args):  
        print(\"WOOT\")  
        print(\"op_target called with args <{0}>\".format(args))  

    def nop():  
       pass  

    new_nop_code = MutableCodeObject(nop.__code__)  
    new_nop_code.co_code = b\"\\x00\" + new_nop_code.co_code[0:3] +  b\"\\x00\" + new_nop_code.co_code[-1:]  
    new_nop_code.co_stacksize += 3  

    nop.__code__ = new_nop_code.get_code()  

    import dis  
    dis.dis(nop)  
    nop()  


    # Don\'t forget that ./python is our custom Python implementing       DEBUG_OP  
    hakril@computer ~/python/CPython3.5 % ./python proof.py  
      8           0 <0>  
                  1 LOAD_CONST               0 (None)  
                  4 <0>  
                  5 RETURN_VALUE  
    WOOT  
    op_target called with args <([], )>  
    WOOT  
    op_target called with args <([None], )>  

看起来它成功了！有一行代码需要说明一下 new_nop_code.co_stacksize += 3

• co_stacksize 表示 code object 所需要的堆栈的大小
• 操作码DEBUG_OP往堆栈中增加了三项，所以我们需要为这些增加的项预留些空间

现在我们可以将我们的操作码注入到每一个 Python 函数中了！

重写字节码

正如我们在上面的例子中所看到的那样，重写 Pyhton 的字节码似乎 so easy。为了在每一个操作码之间注入我们的操作码，我们需要获取每一个操作码的偏移量，然后将我们的操作码注入到这些位置上(把我们操作码注入到参数上是有坏处大大滴)。这些偏移量也很容易获取，使用 dis.Bytecode，就像这样。

    def add_debug_op_everywhere(code_obj):  
         # We get every instruction offset in the code object  
        offsets = [instr.offset for instr in dis.Bytecode(code_obj)]   
        # And insert a DEBUG_OP at every offset  
        return insert_op_debug_list(code_obj, offsets)  

    def insert_op_debug_list(code, offsets):  
         # We insert the DEBUG_OP one by one  
        for nb, off in enumerate(sorted(offsets)):  
            # Need to ajust the offsets by the number of opcodes          already inserted before  
            # That\'s why we sort our offsets!  
            code = insert_op_debug(code, off + nb)  
        return code  

    # Last problem: what does insert_op_debug looks like?  

基于上面的例子，有人可能会想我们的 insert_op_debug 会在指定的偏移量增加一个\”\\x00\”，这尼玛是个坑啊！我们第一个 DEBUG_OP 注入的例子中被注入的函数是没有任何的分支的，为了能够实现完美一个函数注入函数 insert_op_debug 我们需要考虑到存在分支操作码的情况。

Python 的分支一共有两种：

• 绝对分支：看起来是类似这样子的 Instruction_Pointer = argument(instruction)
• 相对分支：看起来是类似这样子的 Instruction_Pointer += argument(instruction)
  • 相对分支总是向前的

我们希望这些分支在我们插入操作码之后仍然能够正常工作，为此我们需要修改一些指令参数。以下是其逻辑流程：

• 对于每一个在插入偏移量之前的相对分支而言
  • 如果目标地址是严格大于我们的插入偏移量的话，将指令参数增加 1
  • 如果相等，则不需要增加 1 就能够在跳转操作和目标地址之间执行我们的操作码DEBUG_OP
  • 如果小于，插入我们的操作码的话并不会影响到跳转操作和目标地址之间的距离
• 对于 code object 中的每一个绝对分支而言
  • 如果目标地址是严格大于我们的插入偏移量的话，将指令参数增加 1
  • 如果相等，那么不需要任何修改，理由和相对分支部分是一样的
  • 如果小于，插入我们的操作码的话并不会影响到跳转操作和目标地址之间的距离

下面是实现：

    # Helper  
    def bytecode_to_string(bytecode):  
        if bytecode.arg is not None:  
            return struct.pack(\" offset:  
                    res_codestring +=      bytecode_to_string(DummyInstr(instr.opcode, instr.arg + 1))  
                    continue  
            res_codestring += bytecode_to_string(instr)  
        # replace_bytecode just replaces the original code co_code  
        return replace_bytecode(code, res_codestring)  

让我们看一下效果如何：

    >>> def lol(x):  
    ...     for i in range(10):  
    ...         if x == i:  
    ...             break  

    >>> dis.dis(lol)  
    101           0 SETUP_LOOP              36 (to 39)  
                  3 LOAD_GLOBAL              0 (range)  
                  6 LOAD_CONST               1 (10)  
                  9 CALL_FUNCTION            1 (1 positional, 0    keyword pair)  
                 12 GET_ITER  
            >>   13 FOR_ITER                22 (to 38)  
                 16 STORE_FAST               1 (i)  

    102          19 LOAD_FAST                0 (x)  
                 22 LOAD_FAST                1 (i)  
                 25 COMPARE_OP               2 (==)  
                 28 POP_JUMP_IF_FALSE       13  

    103          31 BREAK_LOOP  
                 32 JUMP_ABSOLUTE           13  
                 35 JUMP_ABSOLUTE           13  
            >>   38 POP_BLOCK  
            >>   39 LOAD_CONST               0 (None)  
                 42 RETURN_VALUE  
    >>> lol.__code__ = transform_code(lol.__code__,        add_debug_op_everywhere, add_stacksize=3)  


    >>> dis.dis(lol)  
    101           0 <0>  
                  1 SETUP_LOOP              50 (to 54)  
                  4 <0>  
                  5 LOAD_GLOBAL              0 (range)  
                  8 <0>  
                  9 LOAD_CONST               1 (10)  
                 12 <0>  
                 13 CALL_FUNCTION            1 (1 positional, 0   keyword pair)  
                 16 <0>  
                 17 GET_ITER  
            >>   18 <0>  

    102          19 FOR_ITER                30 (to 52)  
                 22 <0>  
                 23 STORE_FAST               1 (i)  
                 26 <0>  
                 27 LOAD_FAST                0 (x)  
                 30 <0>  

    103          31 LOAD_FAST                1 (i)  
                 34 <0>  
                 35 COMPARE_OP               2 (==)  
                 38 <0>  
                 39 POP_JUMP_IF_FALSE       18  
                 42 <0>  
                 43 BREAK_LOOP  
                 44 <0>  
                 45 JUMP_ABSOLUTE           18  
                 48 <0>  
                 49 JUMP_ABSOLUTE           18  
            >>   52 <0>  
                 53 POP_BLOCK  
            >>   54 <0>  
                 55 LOAD_CONST               0 (None)  
                 58 <0>  
                 59 RETURN_VALUE  

     # Setup the simplest handler EVER  
    >>> def op_target(stack, frame):  
    ...     print (stack)  

    # GO  
    >>> lol(2)  
    []  
    []  
    []  
    [10, ]  
    [range(0, 10)]  
    []  
    [0, ]  
    []  
    [2, ]  
    [0, 2, ]  
    [False, ]  
    []  
    [1, ]  
    []  
    [2, ]  
    [1, 2, ]  
    [False, ]  
    []  
    [2, ]  
    []  
    [2, ]  
    [2, 2, ]  
    [True, ]  
    []  
    []  
    [None]  

甚好！现在我们知道了如何获取堆栈信息和 Python 中每一个操作对应的帧信息。上面结果所展示的结果目前而言并不是很实用。在最后一部分中让我们对注入做进一步的封装。

增加 Python 封装

正如您所见到的，所有的底层接口都是好用的。我们最后要做的一件事是让 op_target 更加方便使用(这部分相对而言比较空泛一些，毕竟在我看来这不是整个项目中最有趣的部分)。

首先我们来看一下帧的参数所能提供的信息，如下所示：

• f_code当前帧将执行的 code object
• f_lasti当前的操作(code object 中的字节码字符串的索引)

经过我们的处理我们可以得知 DEBUG_OP 之后要被执行的操作码，这对我们聚合数据并展示是相当有用的。

新建一个用于追踪函数内部机制的类：

• 改变函数自身的 co_code
• 设置回调函数作为 op_debug 的目标函数

一旦我们知道下一个操作，我们就可以分析它并修改它的参数。举例来说我们可以增加一个 auto-follow-called-functions 的特性。

    def op_target(l, f, exc=None):  
        if op_target.callback is not None:  
            op_target.callback(l, f, exc)  

    class Trace:  
        def __init__(self, func):  
            self.func = func  

        def call(self, *args, **kwargs):  
             self.add_func_to_trace(self.func)  
            # Activate Trace callback for the func call  
            op_target.callback = self.callback  
            try:  
                res = self.func(*args, **kwargs)  
            except Exception as e:  
                res = e  
            op_target.callback = None  
            return res  

        def add_func_to_trace(self, f):  
            # Is it code? is it already transformed?  
            if not hasattr(f ,\"op_debug\") and hasattr(f, \"__code__\"):  
                f.__code__ = transform_code(f.__code__,    transform=add_everywhere, add_stacksize=ADD_STACK)  
                f.__globals__[\'op_target\'] = op_target  
                f.op_debug = True  

        def do_auto_follow(self, stack, frame):  
            # Nothing fancy: FrameAnalyser is just the wrapper that  gives the next executed instruction  
            next_instr = FrameAnalyser(frame).next_instr()  
            if \"CALL\" in next_instr.opname:  
                arg = next_instr.arg  
                f_index = (arg & 0xff) + (2 * (arg >> 8))  
                called_func = stack[f_index]  

                # If call target is not traced yet: do it  
                if not hasattr(called_func, \"op_debug\"):  
                    self.add_func_to_trace(called_func)  

现在我们实现一个 Trace 的子类，在这个子类中增加 callback 和 doreport 这两个方法。callback 方法将在每一个操作之后被调用。doreport 方法将我们收集到的信息打印出来。

这是一个伪函数追踪器实现：

    class DummyTrace(Trace):  
        def __init__(self, func):  
            self.func = func  
            self.data = collections.OrderedDict()  
            self.last_frame = None  
            self.known_frame = []  
            self.report = []  

        def callback(self, stack, frame, exc):  
             if frame not in self.known_frame:  
                self.known_frame.append(frame)  
                self.report.append(\" === Entering New Frame {0} ({1})     ===\".format(frame.f_code.co_name, id(frame)))  
                self.last_frame = frame  
            if frame != self.last_frame:  
                self.report.append(\" === Returning to Frame {0}     {1}===\".format(frame.f_code.co_name, id(frame)))  
                self.last_frame = frame  

            self.report.append(str(stack))  
            instr = FrameAnalyser(frame).next_instr()  
            offset = str(instr.offset).rjust(8)  
            opname = str(instr.opname).ljust(20)  
            arg = str(instr.arg).ljust(10)  
            self.report.append(\"{0}  {1} {2} {3}\".format(offset,    opname, arg, instr.argval))  
            self.do_auto_follow(stack, frame)  

        def do_report(self):  
            print(\"\\n\".join(self.report))  

这里有一些实现的例子和使用方法。格式有些不方便观看，毕竟我并不擅长于搞这种对用户友好的报告的事儿。

• 例1自动追踪堆栈信息和已经执行的指令
• 例2上下文管理

递推式构造列表(List Comprehensions)的追踪示例。

• 例3伪追踪器的输出
• 例4输出收集的堆栈信息

总结

这个小项目是一个了解 Python 底层的良好途径，包括解释器的 main loop，Python 实现的 C 代码编程、Python 字节码。通过这个小工具我们可以看到 Python 一些有趣构造函数的字节码行为，例如生成器、上下文管理和递推式构造列表。

这里是这个小项目的完整代码。更进一步的，我们还可以做的是修改我们所追踪的函数的堆栈。我虽然不确定这个是否有用，但是可以肯定是这一过程是相当有趣的。

本文作者系OneAPM工程师编译整理。想阅读更多技术文章，请访问OneAPM官方技术博客。