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

时间:2019-01-25 23:50:13

标签: python

假设我具有有效的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代码作为字符串。我无法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

3 个答案:

答案 0 :(得分:5)

更强大的解决方案是使用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
{p}被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模块的错误/局限性。

答案 1 :(得分:1)

我认为有一个小型解析器是为了尝试并考虑到这种奇怪的异常:

import re

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

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

def test_multiline():
    \"""
    asdasdada
sdadd
    \"""
    pass

def test_comment(
    a #)
):
    return [a,
    # ]
a]

def test_escaped_endline():
    return "asdad \
asdsad \
asdas"

def test_nested():
    return {():[[],
{
}
]
}

def test_strings():
    return '\""" asdasd' + \"""
12asd
12312
"asd2" [
\"""

\"""
def test_fake_def_in_multiline()
\"""
    print(123)
a = "def in_string():"
  def after().
    print("NOPE")

\"""Phew this ain't valid syntax\""" def something(): pass

""".strip()

code_string += '\n'


func_list=[]
func = ''
tab  = ''
brackets = {'(':0, '[':0, '{':0}
close = {')':'(', ']':'[', '}':'{'}
string=''
tab_f=''
c1=''
multiline=False
check=False
for line in code_string.split('\n'):
    tab = re.findall(r'^\s*',line)[0]
    if re.findall(r'^\s*def', line) and not string and not multiline:
        func += line + '\n'
        tab_f = tab
        check=True
    if func:
        if not check:
            if sum(brackets.values()) == 0 and not string and not multiline:
                if len(tab) <= len(tab_f):
                    func_list.append(func)
                    func=''
                    c1=''
                    c2=''
                    continue
            func += line + '\n'
        check = False
    for c0 in line:
        if c0 == '#' and not string and not multiline:
            break
        if c1 != '\\':
            if c0 in ['"', "'"]:
                if c2 == c1 == c0 == '"' and string != "'":
                    multiline = not multiline
                    string = ''
                    continue
                if not multiline:
                    if c0 in string:
                        string = ''
                    else:
                        if not string:
                            string = c0
            if not string and not multiline:
                if c0 in brackets:
                    brackets[c0] += 1
                if c0 in close:
                    b = close[c0]
                    brackets[b] -= 1
        c2=c1
        c1=c0

for f in func_list:
    print('-'*40)
    print(f)

输出:

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

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

----------------------------------------
def baz():
  return [
1,
  ]

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

----------------------------------------
def test_multiline():
    """
    asdasdada
sdadd
    """
    pass

----------------------------------------
def test_comment(
    a #)
):
    return [a,
    # ]
a]

----------------------------------------
def test_escaped_endline():
    return "asdad asdsad asdas"

----------------------------------------
def test_nested():
    return {():[[],
{
}
]
}

----------------------------------------
def test_strings():
    return '""" asdasd' + """
12asd
12312
"asd2" [
"""

----------------------------------------
  def after():
    print("NOPE")

答案 2 :(得分:1)

我会使用python本身,而不是重新发明解析器。

基本上,我会使用compile()内置函数,该函数可以通过编译来检查字符串是否为有效的python代码。我将由选定行组成的字符串传递给它,从每个def到不会编译失败的较远行。

code_string = """
#A comment
def foo(a, b):
  return a + b

def bir(a, b):
  c = a + b
  return c

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

def baz():
  return [
1,
  ]

""".strip()

lines = code_string.split('\n')

#looking for lines with 'def' keywords
defidxs = [e[0] for e in enumerate(lines) if 'def' in e[1]]

#getting the indentation of each 'def'
indents = {}
for i in defidxs:
    ll = lines[i].split('def')
    indents[i] = len(ll[0])

#extracting the strings
end = len(lines)-1
while end > 0:
    if end < defidxs[-1]:
        defidxs.pop()
    try:
        start = defidxs[-1]
    except IndexError: #break if there are no more 'def'
        break

    #empty lines between functions will cause an error, let's remove them
    if len(lines[end].strip()) == 0:
        end = end -1
        continue

    try:
        #fix lines removing indentation or compile will not compile
        fixlines = [ll[indents[start]:] for ll in lines[start:end+1]] #remove indentation
        body = '\n'.join(fixlines)
        compile(body, '<string>', 'exec') #if it fails, throws an exception
        print(body)
        end = start #no need to parse less line if it succeed.
    except:
        pass

    end = end -1

由于except子句没有特定的异常,这有点讨厌,通常不建议这样做,但是没有办法知道可能导致compile失败的原因,所以我不知道如何避免。

这将打印

def baz():
  return [
1,
  ]
def __init__(self):
  self.my_list = [
      'a',
      'b',
  ]
def bir(a, b):
  c = a + b
  return c
def foo(a, b):
  return a + b

请注意,这些功能的打印顺序与code_strings

中显示的功能相反

这甚至可以处理怪异的缩进代码,但是我认为如果您嵌套了函数,它将失败。