用AST重写代码；Python

Question

我正在学习AST，它似乎是一个强大的工具，但我困惑的是代码去哪了以及为什么会消失。如果我想重写代码，该怎么办呢？

example = """def fake(x):\n
    y = ['useless list']\n
    return x
"""

as

example = """def fake(x):\n
    return x
"""

我看不到以这种方式重写的任何方法。我甚至找不到获取该行文本的方法：

In [1]: example = """def fake(x):\n
   ...:     y = ['useless list']\n
   ...:     return x
   ...: """

In [3]: import ast

In [4]: p = ast.parse(example)

In [5]: p
Out[5]: <_ast.Module at 0x7f22f7274a10>

In [6]: p.body
Out[6]: [<_ast.FunctionDef at 0x7f22f7274a50>]

In [7]: p.body
Out[7]: [<_ast.FunctionDef at 0x7f22f7274a50>]

In [8]: f = p.body[0]

In [9]: f
Out[9]: <_ast.FunctionDef at 0x7f22f7274a50>

In [10]: f.body
Out[10]: [<_ast.Assign at 0x7f22f7274b10>, <_ast.Return at 0x7f22f7274c10>]

In [11]: f.name
Out[11]: 'fake'

In [12]: newf = f.body[1:]

In [13]: newf
Out[13]: [<_ast.Return at 0x7f22f7274c10>]

In [14]: z = newf[0]

In [15]: z.value
Out[15]: <_ast.Name at 0x7f22f7274c50>

In [16]: z.value.id
Out[16]: 'x'

更令人惊讶的是它只提供了起始行号，而没有结束行号。这意味着你知道函数从哪里开始，但不知道它在哪里结束，这毫无用处。

我该如何在不使用列表y的情况下获取代码并重写此函数？谢谢。

- codyc4321

---

相关链接：https://dev59.com/lnRA5IYBdhLWcg3w9ivq - wim

Answer 1

也许这能帮助您:

import ast
import astor

example = """
def fake(x):
    y = ['useless list']
    return x
"""

tree = ast.parse(example)

# iterating through list which is represents function on ast
for ind, item in enumerate(tree.body[0].body):
    if isinstance(item, ast.Assign) and isinstance(item.value, ast.List):
        del tree.body[0].body[ind]
        break

print astor.to_source(tree)

更好地使用面向对象编程：

class RemoveList(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        self.generic_visit(node)
        for leaf in node.body:
            if isinstance(leaf, ast.Assign) and isinstance(leaf.value, ast.List):
                del node.body[node.body.index(leaf)]
        ast.fix_missing_locations(node)
        return node


tree_class = ast.parse(example)
remove_list = RemoveList().visit(tree_class)
print astor.to_source(tree_class)

更好的方法是使用RedBaron完整语法树：

from redbaron import RedBaron

example = """
def fake(x):
    y = ['useless list']
    return x
"""
red = RedBaron(example)
methods = red.find_all('DefNode').data
for data in methods:
    if len(data.value.data) > 0:
        for vals in data.value.data:
            if vals[0].type == 'assignment' and vals[0].value.type == 'list':
                index = vals[0].index_on_parent
                par = vals[0].parent
                del par[index]
                break

print red.dumps()

- memu

---

感谢您的回答。RedBaron是一个非常棒的库，它可以保留所有代码源格式，并且使遍历节点变得更加容易。 - rob111

Answer 2

现在有一个库可以告诉您每个节点的确切文本范围：ASTTokens。以下是您的示例：

import ast, asttokens
atok = asttokens.ASTTokens(example, parse=True)

for node in ast.walk(atok.tree):
  if hasattr(node, 'lineno'):
    print atok.get_text_range(node), node.__class__.__name__, atok.get_text(node)

这将产生以下输出（请注意文本中的位置）：

(0, 50) FunctionDef def fake(x):
    y = ['useless list']
    return x
(17, 37) Assign y = ['useless list']
(42, 50) Return return x
(9, 10) Name x
(17, 18) Name y
(21, 37) List ['useless list']
(49, 50) Name x
(22, 36) Str 'useless list'

- DS.

Answer 3

你可以使用ast本身来反解析表达式树：

parsed = ast.parse(text)

source = ast.unparse(parsed)

- Jhkcia