從源代碼字符串中提取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

我認為一個小的解析器是為了嘗試並考慮這個奇怪的例外：

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

Answer 3

我不會重新發明解析器，而是使用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出現的順序相反

這應該處理甚至奇怪的縮進代碼，但我認為如果你有嵌套函數它會失敗。