获取展示 Python 模块中所有使用过的类、方法和函数

在使用新模块时,了解模块中哪些实体被实际使用过有时会很有帮助。我在博文 “Instrumenting Java Code to Find and Handle Unused Classes “中写过类似的内容,但这次我需要在 Python 中使用,而且是方法级粒度的。

简要说明

从 GitHub 下载 trace.py,用它在错误输出中打印调用树和已用方法与类的列表:

import trace
trace.setup(r"MODULE_REGEX", print_location_=True)

实现

这可能是一个很难解决的问题,但当我们使用 sys.settrace 为每个方法和函数调用设置一个处理程序,问题就不难解决了。

基本上有六种不同类型的函数(示例代码在 GitHub 上):

def log(message: str):
    print(message)


class TestClass:
    # static initializer of the class
    x = 100

    def __init__(self):
        # constructor
        log("instance initializer")

    def instance_method(self):
        # instance method, self is bound to an instance
        log("instance method")

    @staticmethod
    def static_method():
        log("static method")

    @classmethod
    def class_method(cls):
        log("class method")


def free_function():
    log("free function")

这一点很重要,因为在下文中我们必须以不同的方式处理它们。但首先,让我们定义几个助手和配置变量:

indent = 0
module_matcher: str = ".*"
print_location: bool = False

我们还想打印方法调用树,因此使用缩进来跟踪当前的缩进级别。module_matcher 是正则表达式,我们用它来决定是否要考虑一个模块、它的类和方法。例如,可以使用 __main__ 来只考虑主模块。print_location 告诉我们是否要打印调用树中每个元素的路径和行位置。

现在来看主辅助类:

def log(message: str):
    print(message, file=sys.stderr)


STATIC_INIT = "<static init>"

@dataclass
class ClassInfo:
    """ Used methods of a class """
    name: str
    used_methods: Set[str] = field(default_factory=set)

    def print(self, indent_: str):
        log(indent_ + self.name)
        for method in sorted(self.used_methods):
            log(indent_ + "  " + method)

    def has_only_static_init(self) -> bool:
        return (
                    len(self.used_methods) == 1 and
                    self.used_methods.pop() == STATIC_INIT)

used_classes: Dict[str, ClassInfo] = {}
free_functions: Set[str] = set()

ClassInfo 保存了一个类的常用方法。我们将使用过的类的 ClassInfo 实例和自由函数存储在全局变量中。

现在,我们将调用处理程序传递给 sys.settrace

def handler(frame: FrameType, event: str, *args):
    """ Trace handler that prints and tracks called functions """
    # find module name
    module_name: str = mod.__name__ if (
        mod := inspect.getmodule(frame.f_code)) else ""

    # get name of the code object
    func_name = frame.f_code.co_name

    # check that the module matches the define regexp
    if not re.match(module_matcher, module_name):
        return
    
    # keep indent in sync
    # this is the only reason why we need
    # the return events and use an inner trace handler
    global indent
    if event == 'return':
        indent -= 2
        return
    if event != "call":
        return

    # insert the current function/method
    name = insert_class_or_function(module_name, func_name, frame)

    # print the current location if neccessary
    if print_location:
        do_print_location(frame)
    
    # print the current function/method
    log(" " * indent + name)

    # keep the indent in sync
    indent += 2

    # return this as the inner handler to get
    # return events
    return handler


def setup(module_matcher_: str = ".*", print_location_: bool = False):
    # ...
    sys.settrace(handler)

现在,我们 “只 “需要获取代码对象的名称,并将其正确收集到 ClassInfo 实例或自由函数集中。基本情况很简单:当当前frame 包含一个局部变量 self 时,我们可能有一个实例方法;当当前frame 包含一个 cls 变量时,我们有一个类方法。

def insert_class_or_function(module_name: str, func_name: str,
                             frame: FrameType) -> str:
    """ Insert the code object and return the name to print """
    if "self" in frame.f_locals or "cls" in frame.f_locals:
        return insert_class_or_instance_function(module_name,
                                                 func_name, frame)
   # ...

def insert_class_or_instance_function(module_name: str,
                                      func_name: str,
                                      frame: FrameType) -> str:
    """
    Insert the code object of an instance or class function and
    return the name to print
    """
    class_name = ""

    if "self" in frame.f_locals:
        # instance methods
        class_name = frame.f_locals["self"].__class__.__name__

    elif "cls" in frame.f_locals:
        # class method
        class_name = frame.f_locals["cls"].__name__
        # we prefix the class method name with "<class>"
        func_name = "<class>" + func_name
    
    # add the module name to class name
    class_name = module_name + "." + class_name
    get_class_info(class_name).used_methods.add(func_name)
    used_classes[class_name].used_methods.add(func_name)
    
    # return the string to print in the class tree
    return class_name + "." + func_name

那么其他三种情况呢?我们使用方法的header line来区分它们:

class StaticFunctionType(Enum):
    INIT = 1
    """ static init """
    STATIC = 2
    """ static function """
    FREE = 3
    """ free function, not related to a class """


def get_static_type(code: CodeType) -> StaticFunctionType:
    file_lines = Path(code.co_filename).read_text().split("\n")
    line = code.co_firstlineno
    header_line = file_lines[line - 1]
    if "class " in header_line:
        # e.g. "class TestClass"
        return StaticFunctionType.INIT
    if "@staticmethod" in header_line:
        return StaticFunctionType.STATIC
    return StaticFunctionType.FREE

当然,这些只是近似值,但对于用于探索的小型实用程序来说,它们已经足够好用了。

如果你还知道其他不使用 Python AST 的方法,请在下面的评论中留言。

使用 get_static_type 函数,我们现在可以完成 insert_class_or_function 函数了:

def insert_class_or_function(module_name: str, func_name: str,
                             frame: FrameType) -> str:
    """ Insert the code object and return the name to print """
    if "self" in frame.f_locals or "cls" in frame.f_locals:
        return insert_class_or_instance_function(module_name,
                                                 func_name, frame)
    # get the type of the current code object
    t = get_static_type(frame.f_code)

    if t == StaticFunctionType.INIT:
        # static initializer, the top level class code
        # func_name is actually the class name here,
        # but classes are technically also callable function
        # objects
        class_name = module_name + "." + func_name
        get_class_info(class_name).used_methods.add(STATIC_INIT)
        return class_name + "." + STATIC_INIT
    
    elif t == StaticFunctionType.STATIC:
        # @staticmethod
        # the qualname is in our example TestClass.static_method,
        # so we have to drop the last part of the name to get
        # the class name
        class_name = module_name + "." + frame.f_code.co_qualname[
                                         :-len(func_name) - 1]
        # we prefix static class names with "<static>"
        func_name = "<static>" + func_name
        get_class_info(class_name).used_methods.add(func_name)
        return class_name + "." + func_name
 
    free_functions.add(frame.f_code.co_name)
    return module_name + "." + func_name

最后要做的是注册一个teardown处理程序,以便在退出时打印收集到的信息:

def teardown():
    """ Teardown the tracer and print the results """
    sys.settrace(None)
    log("********** Trace Results **********")
    print_info()


# trigger teardown on exit
atexit.register(teardown)

使用方法

现在,我们在示例程序的开头加上前缀

import trace

trace.setup(r"__main__")

收集 __main__ 模块的所有信息,并直接传递给 Python 解释器。

我们在程序中添加一些代码来调用所有方法/函数:

def all_methods():
    log("all methods")
    TestClass().instance_method()
    TestClass.static_method()
    TestClass.class_method()
    free_function()


all_methods()

我们的实用程序库在执行时会打印出以下内容:

standard error:

    __main__.TestClass.<static init>
    __main__.all_methods
      __main__.log
      __main__.TestClass.__init__
        __main__.log
      __main__.TestClass.instance_method
        __main__.log
      __main__.TestClass.<static>static_method
        __main__.log
      __main__.TestClass.<class>class_method
        __main__.log
      __main__.free_function
        __main__.log
    ********** Trace Results **********
    Used classes:
      only static init:
      not only static init:
       __main__.TestClass
         <class>class_method
         <static init>
         <static>static_method
         __init__
         instance_method
    Free functions:
      all_methods
      free_function
      log

standard output:

    all methods
    instance initializer
    instance method
    static method
    class method
    free function

结论

这个小工具利用 sys.settrace(和一些字符串处理)的强大功能来查找模块使用的类、方法和函数以及调用树。在试图掌握模块的内部结构和自己的应用程序代码中转使用的模块实体时,该工具非常有用。

我在 GitHub 上以 MIT 许可发布了这段代码,所以请随意改进、扩展和修改它。过几周再来看看我为什么要开发这个工具…

本文文字及图片出自 Finding all used Classes, Methods and Functions of a Python Module

阅读余下内容
 

发表回复

您的电子邮箱地址不会被公开。 必填项已用 * 标注


京ICP备12002735号