bing py代码重构工具

AI_astcode.py

from ipywidgets import interact
import ipywidgets as widgets
from lib2to3.refactor import RefactoringTool,get_all_fix_names as lsFix
import ast,astor,pickle

@interact
def convert(code=widgets.Textarea(value="map({it+1},[0])|{[*it]}, 1|{it+1}", description='Input Code'), fixes=widgets.SelectMultiple(options=lsFix('lib2to3.fixes'), description='Fixes')):    
    if not fixes:
        tree = ast.parse(code)
        # 使用自定义的访问器类遍历并修改语法树
        tree = RefactorVisitor().visit(tree)
        print(s:=astor.to_source(tree), astor.dump_tree(tree.body[0]))
        try: return eval(s)
        except:pass
    tool = RefactoringTool(['lib2to3.fixes.fix_'+k for k in fixes])
    try: r=tool.refactor_string(code+'\n', 'input.py'); print(r);return r
    except:pass

    
# 定义一个自定义的节点访问器类，用于替换表达式
class RefactorVisitor(ast.NodeTransformer):
  def visit_Set(self, node):
    # 如果是花括号包围的表达式
    match node.elts:
      case [expr]:
        return ast.Lambda(args=ast.arguments(args=[ast.arg(arg="it")], defaults=[]), body=expr) #kwonlyargs=[], kw_defaults=[]
      case _: return node

  def visit_BinOp(self, node):
    # 如果是 '|'，右边是花括号包围的表达式
    match (node.op, node.left, node.right):
      case (ast.BitOr(), num, ast.Set(elts=[expr])):
        # 创建一个赋值表达式节点（walrus operator）
        it = ast.NamedExpr(target=ast.Name(id="it", ctx=ast.Store()), value=self.visit(num))
        # 创建一个元组节点，包含赋值表达式和表达式，并用切片操作取最后一个元素（-1）
        body = ast.Tuple(elts=[it, expr], ctx=ast.Load())
        slice = ast.Index(value=ast.Constant(value=-1))
        return ast.Subscript(value=body, slice=slice, ctx=ast.Load())
      case _:return node

        
evals=lambda x:x
class ForLoopInlines(ast.NodeTransformer):
    def __init__(o):o.k=''; o.fn={}; o.kr={}
    def visit(o, node):
        vi=lambda e:[o.visit(pickle.loads(pickle.dumps(e))) for e in e]
        incy=lambda t,k:k+str(t.__setitem__(k,r:=t.get(k,0)+1 ) or r)
        # 如果for循环中有evals函数
        match node:
            case ast.Name(k):return ast.Name(k) if k!=o.k else o.v
            case ast.For(k, ast.Call(f), body) if f.id == "evals":
                # 对参数进行求值，并转换成列表
                o.k=k.id; values = list(eval(astor.to_source(node.iter)))
                new_nodes = []
                # 遍历列表中的每个值
                for x in values:
                    # 复制原始循环体，并替换其中所有出现循环变量的地方为常量
                    o.v = ast.copy_location(ast.Constant(x), k)
                    # 将赋值语句和复制后的循环体添加到新节点列表中
                    new_nodes.extend(vi(body))
                # 返回新节点列表作为替换结果
                return new_nodes
            case ast.Call(f) if (e:=o.fn.get(f.id)):
                # 创建一个新的赋值语句节点，将参数赋值给var
                assign = [ast.Assign(targets=[ast.Name(k.arg)], value=o.visit(node.args[i])) for i,k in enumerate(e[0].args)]
                o.k=e[2]; o.v=ast.Name(incy(o.kr,e[2]))
                # 调用=>语句列表，包含赋值和返回
                f.i.extend(vi([*assign, *e[1] ]))
                return o.v
        return super().visit(node)
    def pos(o,e):
        for e in ast.walk(e):
            match e:
                case ast.FunctionDef(k,ka, body,[deco]) if deco.id=='evals': o.fn[k]=(ka,body[0:-1], body[-1].value.id)
                case ast.Expr(ast.Call(f)) if f.id in o.fn:
                    a,i=e.i; a.insert(i,ast.Module(body:=[])) #f(f(0)) ins pos
                case ast.Call(f) if f.id in o.fn: f.i=body
                case ast.Module(st):
                    for i,s in enumerate(st): s.i=(e.body,i)


# 定义原始代码字符串
code = "f=print\nfor x in evals([1+0,2,3]): f(x+1)"
# 将代码字符串转换成AST对象
tree = ast.parse(code)
tree =ast.parse('@evals\ndef f(x): x+=1; x+=2; return x \nf(f(0))')
# 创建ForLoopExpander对象并应用到AST对象上
expander = ForLoopInlines()
expander.pos(tree)
tree = expander.visit(tree)
# 将修改后的AST对象转换回代码字符串并打印输出
new_code = astor.to_source(tree)
print(new_code)