如何在Python文件中的函数前添加一行

时间:2018-12-26 17:35:48

标签: python abstract-syntax-tree inspect

我正在开发一个工具,该工具需要将内容写入Python文件中的函数/方法的开头。我正在使用Python的inspect从文件中获取函数。给定这些对象,我正在使用inspect.getsourcelines(function_obj)来获取函数的内容。从那里,我的解决方案是解析函数头,然后在此之后写入必要的内容。我的正则表达式是:

re.compile('def \w*\((\w|\s|,|=|\r|\*|\n|^\))*\):')

这感觉非常脆弱,我很肯定它缺少边缘保护套。还有更优雅/更坚固的东西吗?这对ast模块来说是个好用例

1 个答案:

答案 0 :(得分:0)

我不确定为什么我的答案被删除了,但是下面的两个函数是我将行语句添加到Python文件中的函数开头的问题的解决方案。

我使用术语“例程”来表示文件中的任何函数或方法。请注意,从AST节点获取文件似乎并不容易,因此我将其与AST节点一起存储。

import ast

def get_routines_from_file(repo_file):
    '''
    Returns the methods and functions from the file
    '''
    routines = []

    with open(repo_file) as file:
        repo_file_content = file.read()
        repo_module = ast.parse(repo_file_content)
        for node in ast.walk(repo_module):
            if isinstance(node, ast.FunctionDef):
                routines.append((node, repo_file))

    return routines


def prepend_statement(self, file, node, statement):
    '''
    Writes a statement to the beginning of the routine
    '''
    first_node = node.body[0]
    first_line_in_routine = first_node.lineno

    # you unfortunately can't used first_node.col_offset because comments
    # return -1 as their col_offset.
    column_offset = node.col_offset + 4
    indentation = ' ' * column_offset
    statement = indentation + statement + '\n'

    with open(file, 'r') as f:
        contents = f.readlines()

    contents.insert(first_line_in_routine - 1, statement)

    with open(file, 'w') as f:
        f.writelines(contents)