python - 从源代码字符串中提取 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 的方法，因为在 __init__ 这样的情况下，AST 基本上没有足够的信息。

我不能使用 inspect 模块，因为它只适用于“事件对象”，而且我只有 Python 代码作为字符串。我无法评估代码，因为这是一个巨大的安全问题。

理论上我可以为 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 模块。以下代码可以处理奇怪的缩进、注释、多行标记、单行功能 block 和功能 block 内的空行:

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')

可以看到continuation line字面意思是一行'a = 1\n'，只有一个行号25。不幸的是，这显然是 tokenize 模块的错误/限制。