从源代码字符串中提取Python函数源文本

Question

假设我有一个有效的Python源代码,作为一个字符串:

code_string = """
# A comment.
def foo(a, b):
  return a + b
class Bar(object):
  def __init__(self):
    self.my_list = [
        'a',
        'b',
    ]
""".strip()

目标:我想获取包含函数定义源代码的行,保留空格.对于上面的代码字符串,我想得到字符串

def foo(a, b):
  return a + b

和

  def __init__(self):
    self.my_list = [
        'a',
        'b',
    ]

或者,相当于,我很乐意在代码字符串中获取函数的行号:foo跨越2-3行,并__init__跨越5-9行.

尝试

我可以将代码字符串解析为AST:

code_ast = ast.parse(code_string)

我可以找到FunctionDef节点,例如:

function_def_nodes = [node for node in ast.walk(code_ast)
                      if isinstance(node, ast.FunctionDef)]

每个FunctionDef节点的lineno属性告诉我们该函数的第一行.我们可以用以下方法估计该函数的最后一行:

last_line = max(node.lineno for node in ast.walk(function_def_node)
                if hasattr(node, 'lineno'))

但是当函数与不显示为AST节点语法元素结束,例如过去的这不很好地工作]在__init__.

我怀疑有一种只使用AST的方法,因为AST基本上没有足够的信息__init__.

我不能使用该inspect模块,因为它只适用于"活动对象",我只将Python代码作为字符串.我不能eval代码,因为这是一个巨大的安全问题.

从理论上讲,我可以为Python编写一个解析器,但这看起来真的有些过分.

注释中建议的启发式是使用行的前导空格.但是,这可以打破奇怪但有效的函数与奇怪的缩进,如:

def baz():
  return [
1,
  ]

class Baz(object):
  def hello(self, x):
    return self.hello(
x - 1)

def my_type_annotated_function(
  my_long_argument_name: SomeLongArgumentTypeName
) -> SomeLongReturnTypeName:
  # This function's indentation isn't unusual at all.
  pass

Answer 1

一个更强大的解决方案是使用该tokenize模块.以下代码可以处理奇怪的缩进,注释,多行令牌,单行功能块和功能块中的空行:

import tokenize
from io import BytesIO
from collections import deque
code_string = """
# A comment.
def foo(a, b):
  return a + b

class Bar(object):
  def __init__(self):

    self.my_list = [
        'a',
        'b',
    ]

  def test(self): pass
  def abc(self):
    '''multi-
    line token'''

def baz():
  return [
1,
  ]

class Baz(object):
  def hello(self, x):
    a = \
1
    return self.hello(
x - 1)

def my_type_annotated_function(
  my_long_argument_name: SomeLongArgumentTypeName
) -> SomeLongReturnTypeName:
  pass
  # unmatched parenthesis: (
""".strip()
file = BytesIO(code_string.encode())
tokens = deque(tokenize.tokenize(file.readline))
lines = []
while tokens:
    token = tokens.popleft()
    if token.type == tokenize.NAME and token.string == 'def':
        start_line, _ = token.start
        last_token = token
        while tokens:
            token = tokens.popleft()
            if token.type == tokenize.NEWLINE:
                break
            last_token = token
        if last_token.type == tokenize.OP and last_token.string == ':':
            indents = 0
            while tokens:
                token = tokens.popleft()
                if token.type == tokenize.NL:
                    continue
                if token.type == tokenize.INDENT:
                    indents += 1
                elif token.type == tokenize.DEDENT:
                    indents -= 1
                    if not indents:
                        break
                else:
                    last_token = token
        lines.append((start_line, last_token.end[0]))
print(lines)

这输出:

[(2, 3), (6, 11), (13, 13), (14, 16), (18, 21), (24, 27), (29, 33)]

但请注意延续线:

a = \
1

tokenize即使它实际上是两行,也被视为一行,因为如果你打印令牌:

TokenInfo(type=53 (OP), string=':', start=(24, 20), end=(24, 21), line='  def hello(self, x):\n')
TokenInfo(type=4 (NEWLINE), string='\n', start=(24, 21), end=(24, 22), line='  def hello(self, x):\n')
TokenInfo(type=5 (INDENT), string='    ', start=(25, 0), end=(25, 4), line='    a = 1\n')
TokenInfo(type=1 (NAME), string='a', start=(25, 4), end=(25, 5), line='    a = 1\n')
TokenInfo(type=53 (OP), string='=', start=(25, 6), end=(25, 7), line='    a = 1\n')
TokenInfo(type=2 (NUMBER), string='1', start=(25, 8), end=(25, 9), line='    a = 1\n')
TokenInfo(type=4 (NEWLINE), string='\n', start=(25, 9), end=(25, 10), line='    a = 1\n')
TokenInfo(type=1 (NAME), string='return', start=(26, 4), end=(26, 10), line='    return self.hello(\n')

你可以看到延续线实际上被视为一行' a = 1\n',只有一个行号25.tokenize遗憾的是,这显然是模块的错误/限制.

Answer 2

从 Python 3.8 开始，ast.FunctionDef对象（以及所有其他 AST 节点对象）具有该end_lineno字段，因此不再需要复杂的解决方案或脆弱的启发式方法。

节点跨越从 1到1FunctionDef的f线。不包含函数装饰器，应使用 AST 单独处理。f.linenof.end_lineno