神奇的python追踪术---sys.setprofile

python的sys.setprofile 可以设置一个profiling函数,这个profiling函数会在每一个可调用对象被调用和结束时各执行一次。

最近工作中在对superset进行二次开发,尽管源码里增加了类型标注,但定位问题追踪问题仍然是一件不容易的事情。困难之处在于虽然定位了程序执行的起点,却不能在源码里按图索骥找到程序逻辑执行的下一处代码,只能通过不断的在源码里添加print来确定程序的调用链。

烦躁之时,萌生了一个想法,可不可以通过技术手段来获取程序执行过程中的每一次函数调用信息呢,获得函数的名称,所在的文件,这样就能快速的查看整个调用链路上的代码了。

1. sys._getframe

首先想到的是sys._getframe(), 这个函数获得当前的调用堆栈信息,由此向上回溯,即可得到详细的调用信息,这个方法是可行的,也是我最初所采用的方法。但该方法并不能完全满足我的要求,我期望的是确定程序的起点和终点,通过技术手段获得中间过程中的调用信息,起点和终点通常是在同一个函数里。而sys._getframe()只能够做到从调用链的最深处获得调用信息,换言之,sys._getframe()只能获得之前函数调用执行的信息,不能获得之后的函数调用信息。

2. trace

经过一番搜索,找到了trace 模块,这个模块可以获取程序执行过程中的详细信息,使用方法也很简单

python -m trace -t py脚本

输出的信息会比较详细,但也正因为比较详细,我没有采用,因为输出太多了,甄别筛选就成了问题。

3. sys.setprofile

最后在stackflow中找到了自己想要的答案,使用sys.setprofile设置一个profiling函数,神奇之处在于,任何可调用对象被调用时,都会触发这个profiling函数的执行,退出时也是如此。

profiling函数必须有三个参数,分别是frame(当前调用栈), event(事件), 和arg(取决于事件类型)

event的值如果是call, 则表示被调用,如果为return 则表示调用结束。frame是函数视角的调用堆栈,通过frame的f_back属性则可以获得函数调用前的堆栈信息,找到调用者。

import sys

EXCLUSIONS = {'<'}          # Ignore <listcomp>, etc. in the function name.
CALL_EVENT_LIST = []


class FileFilter():
    @staticmethod
    def filter(filename):
        return True

    @staticmethod
    def old_filter(filename):
        return True


def tracefunc(frame, event, arg):
    if event == "call":
        tracefunc.stack_level += 1

        unique_id = frame.f_code.co_filename+str(frame.f_lineno)
        if unique_id in tracefunc.memorized:
            return

        # Part of filename MUST be in white list.
        if 'self' in frame.f_locals:
            class_name = frame.f_locals['self'].__class__.__name__
            func_name = class_name + '.' + frame.f_code.co_name
        else:
            func_name = frame.f_code.co_name

        func_name = '{indent}{name}'.format(
                indent=tracefunc.stack_level*2*'-', name=func_name)

        if not FileFilter.filter(frame.f_code.co_filename):
            return

        frame_back = frame.f_back    # 获取调用函数时的信息
        txt = '{: <40} # {}, {}, {}, {}'.format(
                func_name, frame.f_code.co_filename, frame.f_lineno, frame_back.f_code.co_filename, frame_back.f_lineno)

        CALL_EVENT_LIST.append(txt)

        tracefunc.memorized.add(unique_id)

    elif event == "return":
        tracefunc.stack_level -= 1

def start(filter_func=None):
    if filter_func:
        FileFilter.filter = filter_func

    tracefunc.memorized = set()
    tracefunc.stack_level = 0
    CALL_EVENT_LIST.clear()
    sys.setprofile(tracefunc)


def output():
    for text in CALL_EVENT_LIST:
        print(text)

    CALL_EVENT_LIST.clear()


def stop():
    def do_nothing(frame, event, arg):
        pass

    FileFilter.filter = FileFilter.old_filter
    sys.setprofile(do_nothing)

我编写了 start, stop, outpu三个函数,start负责设置profiling函数,stop负责将profiling函数设置为一个只有pass一条语句的函数,这样,我可以随时随地的开启和关闭,只追中探查指定区域内的代码,output函数负责输出函数的调用信息并清空CALL_EVENT_LIST。

tracefunc会处理每一次函数的调用和退出,包括第三方库和python内置库里的函数,这样会导致搜集的信息过多,为此,我设计一个FileFilter类,在调用start函数时,你可以自己传入一个过滤函数,专门用于过滤搜集到的信息。对于我来说,我只想搜集supertset源码里的函数调用过程,那么我就可以这样来调用

start(filter_func=lambda x: x.find('superset/') != -1)

superset实际使用情况,限于篇幅和大家对它并不了解,我这里就不用它做例子了,写4个简单的函数来演示实际效果

def func1():
    print('ok')

def func2():
    func1()

def func3():
    func2()

def func4():
    func3()

def file_filter(filename):
    if filename.find('calc') != -1:
        return True
    return False

start(filter_func=lambda x: x.find('calc') != -1)
func4()
stop()
output()

实际输出效果

ok
func4                                    # D:/test/calc.py, 83, D:/test/calc.py, 92
--func3                                  # D:/test/calc.py, 80, D:/test/calc.py, 84
----func2                                # D:/test/calc.py, 77, D:/test/calc.py, 81
------func1                              # D:/test/calc.py, 74, D:/test/calc.py, 78
stop                                     # D:/test/calc.py, 66, D:/test/calc.py, 93

扫描关注, 与我技术互动

QQ交流群: 211426309

加入知识星球, 每天收获更多精彩内容

分享日常研究的python技术和遇到的问题及解决方案