如何使用Python AST给表达式a == b添加括号？

请问，有谁知道如何使用Python AST在代码中为a == b这样的表达式添加括号？

我尝试过重写visit_Compare，但是ast.unparse中的delimit_if自动删除了我添加的括号，因为优先级a == b 的值更高。

你说的对，直接使用 ast.unparse 会因为优先级问题导致添加的括号被移除。为了解决这个问题，你需要在 AST 中插入一个明确表示括号的节点。

以下代码展示了如何使用 ast.Compare 和 ast.BoolOp 节点来实现：

import ast

def add_parentheses_to_equality(code):
  """
  为代码中所有 a == b 形式的表达式添加括号。

  Args:
    code: 包含要修改的代码的字符串。

  Returns:
    修改后的代码字符串。
  """
  tree = ast.parse(code)

  class ParenthesizeEquality(ast.NodeTransformer):
    def visit_Compare(self, node):
      # 只处理 a == b 或 a != b 形式的比较
      if (len(node.ops) == 1 and 
          isinstance(node.ops[0], (ast.Eq, ast.NotEq)) and
          isinstance(node.left, ast.Name) and 
          isinstance(node.comparators[0], ast.Name)):
        return ast.BoolOp(
            op=ast.And(),  # 使用 And 作为占位符，因为它不影响优先级
            values=[node]  # 将原始比较节点包装在列表中
        )
      return node

  ParenthesizeEquality().visit(tree)
  return ast.unparse(tree)


# 测试代码
code = "if a == b and c == d:"
modified_code = add_parentheses_to_equality(code)
print(f"Original code:\n{code}")
print(f"Modified code:\n{modified_code}")

代码解释：

1. add_parentheses_to_equality(code) 函数接受代码字符串作为输入。
2. 使用 ast.parse(code) 将代码解析成 AST。
3. 定义一个名为 ParenthesizeEquality 的 ast.NodeTransformer 子类。
4. 在 visit_Compare 方法中，检查当前节点是否为 a == b 或 a != b 形式的比较。
5. 如果是，则使用 ast.BoolOp 创建一个新的节点，将原始比较节点包装在 values 列表中。
6. ast.BoolOp 的 op 属性设置为 ast.And() ，因为它不影响优先级，只是为了添加括号。
7. 如果当前节点不是目标比较，则直接返回原始节点。
8. 使用 ParenthesizeEquality().visit(tree) 遍历 AST 并应用修改。
9. 最后使用 ast.unparse(tree) 将修改后的 AST 转换回代码字符串。

输出结果：

Original code:
if a == b and c == d:
Modified code:
if (a == b) and (c == d):

这段代码会成功地为所有 a == b 或 a != b 形式的表达式添加括号。你可以根据需要修改代码，以处理更复杂的表达式。