以编程方式通过lib2to3转换python代码

时间:2016-01-18 06:44:26

标签: python abstract-syntax-tree

我想通过Python的标准lib的lib2to3将Python模块中出现的所有“some_func(a,b)”转换为“断言a == b”。我写了一个脚本,将源作为输入:

# convert.py
# convert assert_equal(a, b) to assert a == b
from lib2to3 import refactor
refac = refactor.RefactoringTool(['fix_assert_equal'])
result = refac.refactor_string('assert_equal(123, 456)\n', 'assert equal')
print(result)

和单独模块中的实际修复程序:

# fix_assert_equal.py, in same folder as convert.py
from lib2to3 import fixer_base, pygram, pytree, pgen2
import ast
import logging

grammar = pygram.python_grammar
logger = logging.getLogger("RefactoringTool")
driver = pgen2.driver.Driver(grammar, convert=pytree.convert, logger=logger)
dest_tree = driver.parse_string('assert a == b\n')


class FixAssertEqual(fixer_base.BaseFix):
    BM_compatible = True

    PATTERN = """
    power< 'assert_equal'
        trailer<
            '('
            arglist<
                obj1=any ','
                obj2=any
            >
            ')'
        >
    >
    """

    def transform(self, node, results):
        assert results
        obj1 = results["obj1"]
        obj1 = obj1.clone()
        obj1.prefix = ""
        obj2 = results["obj2"]
        obj2 = obj2.clone()
        obj2.prefix = ""
        prefix = node.prefix

        dest_tree2 = dest_tree.clone()
        node = dest_tree2.children[0].children[0].children[1]
        node.children[0] = obj1
        node.children[2] = obj2
        dest_tree2.prefix = prefix

        return dest_tree2

然而,这会产生输出assert123 ==456而不是assert 123 == 456。知道如何解决这个问题吗?

1 个答案:

答案 0 :(得分:0)

我找到了解决方案:obj2.prefix不应设置为“”,这会删除对象之前的空格。