调用ast.get_source_segment报错TypeError：Module对象无len()方法求助

问题代码 #

尝试解析包含if语句的Python代码并提取对应源码片段，代码如下：

import ast

tree = ast.parse("""

class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age

    def age_change(self):
        if(self.age>20):
            return self.age + 5
        else:
            return self.age

        """)

for node in ast.walk(tree):
    if(isinstance(node, ast.If)):
        print(ast.get_source_segment(tree, node))

报错信息 #

运行后触发TypeError：

Traceback (most recent call last):
  File "C:\Users\thoma\AppData\Roaming\JetBrains\PyCharmCE2020.3\scratches\scratch_21.py", line 20, in <module>
    print(ast.get_source_segment(tree, node))
  File "C:\Users\thoma\anaconda3\envs\torch_1\lib\ast.py", line 354, in get_source_segment
    lines = _splitlines_no_ff(source)
  File "C:\Users\thoma\anaconda3\envs\torch_1\lib\ast.py", line 307, in _splitlines_no_ff
    while idx < len(source):
TypeError: object of type 'Module' has no len()

报错原因 #

ast.get_source_segment()的第一个参数要求传入原始源代码字符串，而非ast.parse()返回的Module对象。代码中错误传入了tree（即Module实例），函数内部尝试对其调用len()，但Module类型不支持该操作，因此抛出错误。

解决方案 #

将原始源代码保存为独立变量，同时用于ast.parse()和ast.get_source_segment()的第一个参数。

修正后的代码 #

import ast

# 保存原始源代码字符串
source_code = """

class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age

    def age_change(self):
        if(self.age>20):
            return self.age + 5
        else:
            return self.age

        """

tree = ast.parse(source_code)

for node in ast.walk(tree):
    if isinstance(node, ast.If):
        # 传入原始源代码字符串而非Module对象
        print(ast.get_source_segment(source_code, node))

运行后会正确输出if语句的源码片段：

if(self.age>20):
    return self.age + 5
else:
    return self.age

内容的提问来源于stack exchange，提问作者desert_ranger