Add AST mapping, simplify Located Nodes

ast/asdl_rs.py

@@ -249,6 +249,78 @@ class StructVisitor(EmitVisitor):
         else:
             return "", ""
 
+class MapAstImplVisitor(EmitVisitor):
+    def __init__(self, file, typeinfo):
+        self.typeinfo = typeinfo
+        super().__init__(file)
+
+    def get_generics(self, typ):
+        if self.typeinfo[typ].has_userdata:
+            return "<T>", "<U>"
+        else:
+            return "", ""
+
+    def visitModule(self, mod):
+        for dfn in mod.dfns:
+            self.visit(dfn)
+
+    def visitType(self, type, depth=0):
+        self.visit(type.value, type.name, depth)
+
+    def visitSum(self, sum, name, depth):
+        apply_t, apply_u = self.get_generics(name)
+        enumname = get_rust_type(name)
+        if sum.attributes:
+            enumname += "Kind"
+
+        if not apply_t:
+            self.emit(f"no_user!({enumname});", depth)
+            return
+
+
+        self.emit(f"impl<T, U> MapAst<T, U> for {enumname}{apply_t} {{", depth)
+        self.emit(f"type Mapped = {enumname}{apply_u};", depth + 1)
+        self.emit("fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {", depth + 1)
+        self.emit("match self {", depth + 2)
+        for cons in sum.types:
+            fields_pattern = self.make_pattern(cons.fields)
+            self.emit(f"{enumname}::{cons.name} {{ {fields_pattern} }} => {{", depth + 3)
+            self.gen_construction(f"{enumname}::{cons.name}", cons.fields, depth + 4)
+            self.emit("}", depth + 3)
+        self.emit("}", depth + 2)
+        self.emit("}", depth + 1)
+        self.emit("}", depth)
+
+    def visitProduct(self, product, name, depth):
+        apply_t, apply_u = self.get_generics(name)
+        structname = get_rust_type(name)
+        if product.attributes:
+            structname += "Data"
+
+        if not apply_t:
+            self.emit(f"no_user!({structname});", depth)
+            return
+
+        self.emit(f"impl<T, U> MapAst<T, U> for {structname}{apply_t} {{", depth)
+        self.emit(f"type Mapped = {structname}{apply_u};", depth + 1)
+        self.emit("fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {", depth + 1)
+        fields_pattern = self.make_pattern(product.fields)
+        self.emit(f"let {structname} {{ {fields_pattern} }} = self;", depth + 2)
+        self.gen_construction(structname, product.fields, depth + 2)
+        self.emit("}", depth + 1)
+        self.emit("}", depth)
+
+    def make_pattern(self, fields):
+        return ",".join(rust_field(f.name) for f in fields)
+
+    def gen_construction(self, cons_path, fields, depth):
+        self.emit(f"Ok({cons_path} {{", depth)
+        for field in fields:
+            name = rust_field(field.name)
+            self.emit(f"{name}: {name}.try_map_ast(f)?,", depth + 1)
+        self.emit("})", depth)
+
+
 class ClassDefVisitor(EmitVisitor):
 
     def visitModule(self, mod):
@@ -321,21 +393,20 @@ class TraitImplVisitor(EmitVisitor):
         self.visit(type.value, type.name, depth)
 
     def visitSum(self, sum, name, depth):
-        rustname = enumname = get_rust_type(name)
-        node = "self"
+        enumname = get_rust_type(name)
         if sum.attributes:
-            enumname = rustname + "Kind"
-            node = "self.node"
+            enumname += "Kind"
 
-        self.emit(f"impl Node for ast::{rustname} {{", depth)
+
+        self.emit(f"impl NamedNode for ast::{enumname} {{", depth)
+        self.emit(f"const NAME: &'static str = {json.dumps(name)};", depth + 1)
+        self.emit("}", depth)
+        self.emit(f"impl Node for ast::{enumname} {{", depth)
         self.emit("fn ast_to_object(self, _vm: &VirtualMachine) -> PyObjectRef {", depth + 1)
-        self.emit(f"let node = match {node} {{", depth + 2)
+        self.emit("match self {", depth + 2)
         for variant in sum.types:
             self.constructor_to_object(variant, enumname, depth + 3)
-        self.emit("};", depth + 2)
-        if sum.attributes:
-            self.add_location(depth + 2)
-        self.emit("node.into_object()", depth + 2)
+        self.emit("}", depth + 2)
         self.emit("}", depth + 1)
         self.emit("fn ast_from_object(_vm: &VirtualMachine, _object: PyObjectRef) -> PyResult<Self> {", depth + 1)
         self.gen_sum_fromobj(sum, name, enumname, depth + 2)
@@ -349,22 +420,18 @@ class TraitImplVisitor(EmitVisitor):
         self.emit("}", depth)
 
     def visitProduct(self, product, name, depth):
-        rustname = structname = get_rust_type(name)
-        node = "self"
+        structname = get_rust_type(name)
         if product.attributes:
-            structname = rustname + "Data"
-            node = "self.node"
+            structname += "Data"
 
-        self.emit(f"impl Node for ast::{rustname} {{", depth)
+        self.emit(f"impl NamedNode for ast::{structname} {{", depth)
+        self.emit(f"const NAME: &'static str = {json.dumps(name)};", depth + 1)
+        self.emit("}", depth)
+        self.emit(f"impl Node for ast::{structname} {{", depth)
         self.emit("fn ast_to_object(self, _vm: &VirtualMachine) -> PyObjectRef {", depth + 1)
         fields_pattern = self.make_pattern(product.fields)
-        self.emit("let node = {", depth + 2)
-        self.emit(f"let ast::{structname} {{ {fields_pattern} }} = {node};", depth + 3)
-        self.make_node(name, product.fields, depth + 3)
-        self.emit("};", depth + 2)
-        if product.attributes:
-            self.add_location(depth + 2)
-        self.emit("node.into_object()", depth + 2)
+        self.emit(f"let ast::{structname} {{ {fields_pattern} }} = self;", depth + 2)
+        self.make_node(name, product.fields, depth + 2)
         self.emit("}", depth + 1)
         self.emit("fn ast_from_object(_vm: &VirtualMachine, _object: PyObjectRef) -> PyResult<Self> {", depth + 1)
         self.gen_product_fromobj(product, name, structname, depth + 2)
@@ -378,20 +445,17 @@ class TraitImplVisitor(EmitVisitor):
             self.emit("let _dict = _node.as_object().dict().unwrap();", depth)
         for f in fields:
             self.emit(f"_dict.set_item({json.dumps(f.name)}, {rust_field(f.name)}.ast_to_object(_vm), _vm).unwrap();", depth)
-        self.emit("_node", depth)
+        self.emit("_node.into_object()", depth)
 
     def make_pattern(self, fields):
         return ",".join(rust_field(f.name) for f in fields)
 
-    def add_location(self, depth):
-        self.emit(f"node_add_location(&node, self.location, _vm);", depth)
-
     def gen_sum_fromobj(self, sum, sumname, enumname, depth):
         if sum.attributes:
             self.extract_location(sumname, depth)
 
         self.emit("let _cls = _object.class();", depth)
-        self.emit("let node =", depth)
+        self.emit("Ok(", depth)
         for cons in sum.types:
             self.emit(f"if _cls.is(Node{cons.name}::static_type()) {{", depth)
             self.gen_construction(f"{enumname}::{cons.name}", cons, sumname, depth + 1)
@@ -400,25 +464,15 @@ class TraitImplVisitor(EmitVisitor):
         self.emit("{", depth)
         msg = f'format!("expected some sort of {sumname}, but got {{}}",_vm.to_repr(&_object)?)'
         self.emit(f"return Err(_vm.new_type_error({msg}));", depth + 1)
-        self.emit("};", depth)
-
-        if sum.attributes:
-            self.wrap_located_node(depth)
-
-        self.emit("Ok(node)", depth)
+        self.emit("})", depth)
 
     def gen_product_fromobj(self, product, prodname, structname, depth):
         if product.attributes:
             self.extract_location(prodname, depth)
 
-        self.emit("let node =", depth)
+        self.emit("Ok(", depth)
         self.gen_construction(structname, product, prodname, depth + 1)
-        self.emit(";", depth)
-
-        if product.attributes:
-            self.wrap_located_node(depth)
-
-        self.emit("Ok(node)", depth)
+        self.emit(")", depth)
 
     def gen_construction(self, cons_path, cons, name, depth):
         self.emit(f"ast::{cons_path} {{", depth)
@@ -441,9 +495,20 @@ class TraitImplVisitor(EmitVisitor):
         else:
             return f"Node::ast_from_object(_vm, get_node_field(_vm, &_object, {name}, {json.dumps(typename)})?)?"
 
+class ChainOfVisitors:
+    def __init__(self, *visitors):
+        self.visitors = visitors
+
+    def visit(self, object):
+        for v in self.visitors:
+            v.visit(object)
+            v.emit("", 0)
+
+
 def write_ast_def(mod, typeinfo, f):
     f.write('pub use crate::location::Location;\n')
     f.write('pub use crate::constant::*;\n')
+    f.write('use crate::map_ast::MapAst;\n')
     f.write('\n')
     f.write('type Ident = String;\n')
     f.write('\n')
@@ -461,22 +526,19 @@ def write_ast_def(mod, typeinfo, f):
     f.write('}\n')
     f.write('\n')
 
-    StructVisitor(f, typeinfo).visit(mod)
+    c = ChainOfVisitors(StructVisitor(f, typeinfo),
+                        MapAstImplVisitor(f, typeinfo))
+    c.visit(mod)
 
 
 def write_ast_mod(mod, f):
     f.write('use super::*;\n')
     f.write('\n')
 
-    ClassDefVisitor(f).visit(mod)
-
-    f.write('\n')
-
-    TraitImplVisitor(f).visit(mod)
-
-    f.write('\n')
-
-    ExtendModuleVisitor(f).visit(mod)
+    c = ChainOfVisitors(ClassDefVisitor(f),
+                        TraitImplVisitor(f),
+                        ExtendModuleVisitor(f))
+    c.visit(mod)
 
 def main(input_filename, ast_mod_filename, ast_def_filename, dump_module=False):
     auto_gen_msg = AUTOGEN_MESSAGE.format("/".join(Path(__file__).parts[-2:]))

ast/src/ast_gen.rs

@@ -2,6 +2,7 @@
 
 pub use crate::constant::*;
 pub use crate::location::Location;
+use crate::map_ast::MapAst;
 
 type Ident = String;
 
@@ -387,3 +388,437 @@ pub struct Withitem<U = ()> {
 pub enum TypeIgnore {
     TypeIgnore { lineno: usize, tag: String },
 }
+
+impl<T, U> MapAst<T, U> for Mod<T> {
+    type Mapped = Mod<U>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        match self {
+            Mod::Module { body, type_ignores } => Ok(Mod::Module {
+                body: body.try_map_ast(f)?,
+                type_ignores: type_ignores.try_map_ast(f)?,
+            }),
+            Mod::Interactive { body } => Ok(Mod::Interactive {
+                body: body.try_map_ast(f)?,
+            }),
+            Mod::Expression { body } => Ok(Mod::Expression {
+                body: body.try_map_ast(f)?,
+            }),
+            Mod::FunctionType { argtypes, returns } => Ok(Mod::FunctionType {
+                argtypes: argtypes.try_map_ast(f)?,
+                returns: returns.try_map_ast(f)?,
+            }),
+        }
+    }
+}
+impl<T, U> MapAst<T, U> for StmtKind<T> {
+    type Mapped = StmtKind<U>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        match self {
+            StmtKind::FunctionDef {
+                name,
+                args,
+                body,
+                decorator_list,
+                returns,
+                type_comment,
+            } => Ok(StmtKind::FunctionDef {
+                name: name.try_map_ast(f)?,
+                args: args.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+                decorator_list: decorator_list.try_map_ast(f)?,
+                returns: returns.try_map_ast(f)?,
+                type_comment: type_comment.try_map_ast(f)?,
+            }),
+            StmtKind::AsyncFunctionDef {
+                name,
+                args,
+                body,
+                decorator_list,
+                returns,
+                type_comment,
+            } => Ok(StmtKind::AsyncFunctionDef {
+                name: name.try_map_ast(f)?,
+                args: args.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+                decorator_list: decorator_list.try_map_ast(f)?,
+                returns: returns.try_map_ast(f)?,
+                type_comment: type_comment.try_map_ast(f)?,
+            }),
+            StmtKind::ClassDef {
+                name,
+                bases,
+                keywords,
+                body,
+                decorator_list,
+            } => Ok(StmtKind::ClassDef {
+                name: name.try_map_ast(f)?,
+                bases: bases.try_map_ast(f)?,
+                keywords: keywords.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+                decorator_list: decorator_list.try_map_ast(f)?,
+            }),
+            StmtKind::Return { value } => Ok(StmtKind::Return {
+                value: value.try_map_ast(f)?,
+            }),
+            StmtKind::Delete { targets } => Ok(StmtKind::Delete {
+                targets: targets.try_map_ast(f)?,
+            }),
+            StmtKind::Assign {
+                targets,
+                value,
+                type_comment,
+            } => Ok(StmtKind::Assign {
+                targets: targets.try_map_ast(f)?,
+                value: value.try_map_ast(f)?,
+                type_comment: type_comment.try_map_ast(f)?,
+            }),
+            StmtKind::AugAssign { target, op, value } => Ok(StmtKind::AugAssign {
+                target: target.try_map_ast(f)?,
+                op: op.try_map_ast(f)?,
+                value: value.try_map_ast(f)?,
+            }),
+            StmtKind::AnnAssign {
+                target,
+                annotation,
+                value,
+                simple,
+            } => Ok(StmtKind::AnnAssign {
+                target: target.try_map_ast(f)?,
+                annotation: annotation.try_map_ast(f)?,
+                value: value.try_map_ast(f)?,
+                simple: simple.try_map_ast(f)?,
+            }),
+            StmtKind::For {
+                target,
+                iter,
+                body,
+                orelse,
+                type_comment,
+            } => Ok(StmtKind::For {
+                target: target.try_map_ast(f)?,
+                iter: iter.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+                orelse: orelse.try_map_ast(f)?,
+                type_comment: type_comment.try_map_ast(f)?,
+            }),
+            StmtKind::AsyncFor {
+                target,
+                iter,
+                body,
+                orelse,
+                type_comment,
+            } => Ok(StmtKind::AsyncFor {
+                target: target.try_map_ast(f)?,
+                iter: iter.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+                orelse: orelse.try_map_ast(f)?,
+                type_comment: type_comment.try_map_ast(f)?,
+            }),
+            StmtKind::While { test, body, orelse } => Ok(StmtKind::While {
+                test: test.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+                orelse: orelse.try_map_ast(f)?,
+            }),
+            StmtKind::If { test, body, orelse } => Ok(StmtKind::If {
+                test: test.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+                orelse: orelse.try_map_ast(f)?,
+            }),
+            StmtKind::With {
+                items,
+                body,
+                type_comment,
+            } => Ok(StmtKind::With {
+                items: items.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+                type_comment: type_comment.try_map_ast(f)?,
+            }),
+            StmtKind::AsyncWith {
+                items,
+                body,
+                type_comment,
+            } => Ok(StmtKind::AsyncWith {
+                items: items.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+                type_comment: type_comment.try_map_ast(f)?,
+            }),
+            StmtKind::Raise { exc, cause } => Ok(StmtKind::Raise {
+                exc: exc.try_map_ast(f)?,
+                cause: cause.try_map_ast(f)?,
+            }),
+            StmtKind::Try {
+                body,
+                handlers,
+                orelse,
+                finalbody,
+            } => Ok(StmtKind::Try {
+                body: body.try_map_ast(f)?,
+                handlers: handlers.try_map_ast(f)?,
+                orelse: orelse.try_map_ast(f)?,
+                finalbody: finalbody.try_map_ast(f)?,
+            }),
+            StmtKind::Assert { test, msg } => Ok(StmtKind::Assert {
+                test: test.try_map_ast(f)?,
+                msg: msg.try_map_ast(f)?,
+            }),
+            StmtKind::Import { names } => Ok(StmtKind::Import {
+                names: names.try_map_ast(f)?,
+            }),
+            StmtKind::ImportFrom {
+                module,
+                names,
+                level,
+            } => Ok(StmtKind::ImportFrom {
+                module: module.try_map_ast(f)?,
+                names: names.try_map_ast(f)?,
+                level: level.try_map_ast(f)?,
+            }),
+            StmtKind::Global { names } => Ok(StmtKind::Global {
+                names: names.try_map_ast(f)?,
+            }),
+            StmtKind::Nonlocal { names } => Ok(StmtKind::Nonlocal {
+                names: names.try_map_ast(f)?,
+            }),
+            StmtKind::Expr { value } => Ok(StmtKind::Expr {
+                value: value.try_map_ast(f)?,
+            }),
+            StmtKind::Pass {} => Ok(StmtKind::Pass {}),
+            StmtKind::Break {} => Ok(StmtKind::Break {}),
+            StmtKind::Continue {} => Ok(StmtKind::Continue {}),
+        }
+    }
+}
+impl<T, U> MapAst<T, U> for ExprKind<T> {
+    type Mapped = ExprKind<U>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        match self {
+            ExprKind::BoolOp { op, values } => Ok(ExprKind::BoolOp {
+                op: op.try_map_ast(f)?,
+                values: values.try_map_ast(f)?,
+            }),
+            ExprKind::NamedExpr { target, value } => Ok(ExprKind::NamedExpr {
+                target: target.try_map_ast(f)?,
+                value: value.try_map_ast(f)?,
+            }),
+            ExprKind::BinOp { left, op, right } => Ok(ExprKind::BinOp {
+                left: left.try_map_ast(f)?,
+                op: op.try_map_ast(f)?,
+                right: right.try_map_ast(f)?,
+            }),
+            ExprKind::UnaryOp { op, operand } => Ok(ExprKind::UnaryOp {
+                op: op.try_map_ast(f)?,
+                operand: operand.try_map_ast(f)?,
+            }),
+            ExprKind::Lambda { args, body } => Ok(ExprKind::Lambda {
+                args: args.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+            }),
+            ExprKind::IfExp { test, body, orelse } => Ok(ExprKind::IfExp {
+                test: test.try_map_ast(f)?,
+                body: body.try_map_ast(f)?,
+                orelse: orelse.try_map_ast(f)?,
+            }),
+            ExprKind::Dict { keys, values } => Ok(ExprKind::Dict {
+                keys: keys.try_map_ast(f)?,
+                values: values.try_map_ast(f)?,
+            }),
+            ExprKind::Set { elts } => Ok(ExprKind::Set {
+                elts: elts.try_map_ast(f)?,
+            }),
+            ExprKind::ListComp { elt, generators } => Ok(ExprKind::ListComp {
+                elt: elt.try_map_ast(f)?,
+                generators: generators.try_map_ast(f)?,
+            }),
+            ExprKind::SetComp { elt, generators } => Ok(ExprKind::SetComp {
+                elt: elt.try_map_ast(f)?,
+                generators: generators.try_map_ast(f)?,
+            }),
+            ExprKind::DictComp {
+                key,
+                value,
+                generators,
+            } => Ok(ExprKind::DictComp {
+                key: key.try_map_ast(f)?,
+                value: value.try_map_ast(f)?,
+                generators: generators.try_map_ast(f)?,
+            }),
+            ExprKind::GeneratorExp { elt, generators } => Ok(ExprKind::GeneratorExp {
+                elt: elt.try_map_ast(f)?,
+                generators: generators.try_map_ast(f)?,
+            }),
+            ExprKind::Await { value } => Ok(ExprKind::Await {
+                value: value.try_map_ast(f)?,
+            }),
+            ExprKind::Yield { value } => Ok(ExprKind::Yield {
+                value: value.try_map_ast(f)?,
+            }),
+            ExprKind::YieldFrom { value } => Ok(ExprKind::YieldFrom {
+                value: value.try_map_ast(f)?,
+            }),
+            ExprKind::Compare {
+                left,
+                ops,
+                comparators,
+            } => Ok(ExprKind::Compare {
+                left: left.try_map_ast(f)?,
+                ops: ops.try_map_ast(f)?,
+                comparators: comparators.try_map_ast(f)?,
+            }),
+            ExprKind::Call {
+                func,
+                args,
+                keywords,
+            } => Ok(ExprKind::Call {
+                func: func.try_map_ast(f)?,
+                args: args.try_map_ast(f)?,
+                keywords: keywords.try_map_ast(f)?,
+            }),
+            ExprKind::FormattedValue {
+                value,
+                conversion,
+                format_spec,
+            } => Ok(ExprKind::FormattedValue {
+                value: value.try_map_ast(f)?,
+                conversion: conversion.try_map_ast(f)?,
+                format_spec: format_spec.try_map_ast(f)?,
+            }),
+            ExprKind::JoinedStr { values } => Ok(ExprKind::JoinedStr {
+                values: values.try_map_ast(f)?,
+            }),
+            ExprKind::Constant { value, kind } => Ok(ExprKind::Constant {
+                value: value.try_map_ast(f)?,
+                kind: kind.try_map_ast(f)?,
+            }),
+            ExprKind::Attribute { value, attr, ctx } => Ok(ExprKind::Attribute {
+                value: value.try_map_ast(f)?,
+                attr: attr.try_map_ast(f)?,
+                ctx: ctx.try_map_ast(f)?,
+            }),
+            ExprKind::Subscript { value, slice, ctx } => Ok(ExprKind::Subscript {
+                value: value.try_map_ast(f)?,
+                slice: slice.try_map_ast(f)?,
+                ctx: ctx.try_map_ast(f)?,
+            }),
+            ExprKind::Starred { value, ctx } => Ok(ExprKind::Starred {
+                value: value.try_map_ast(f)?,
+                ctx: ctx.try_map_ast(f)?,
+            }),
+            ExprKind::Name { id, ctx } => Ok(ExprKind::Name {
+                id: id.try_map_ast(f)?,
+                ctx: ctx.try_map_ast(f)?,
+            }),
+            ExprKind::List { elts, ctx } => Ok(ExprKind::List {
+                elts: elts.try_map_ast(f)?,
+                ctx: ctx.try_map_ast(f)?,
+            }),
+            ExprKind::Tuple { elts, ctx } => Ok(ExprKind::Tuple {
+                elts: elts.try_map_ast(f)?,
+                ctx: ctx.try_map_ast(f)?,
+            }),
+            ExprKind::Slice { lower, upper, step } => Ok(ExprKind::Slice {
+                lower: lower.try_map_ast(f)?,
+                upper: upper.try_map_ast(f)?,
+                step: step.try_map_ast(f)?,
+            }),
+        }
+    }
+}
+no_user!(ExprContext);
+no_user!(Boolop);
+no_user!(Operator);
+no_user!(Unaryop);
+no_user!(Cmpop);
+impl<T, U> MapAst<T, U> for Comprehension<T> {
+    type Mapped = Comprehension<U>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        let Comprehension {
+            target,
+            iter,
+            ifs,
+            is_async,
+        } = self;
+        Ok(Comprehension {
+            target: target.try_map_ast(f)?,
+            iter: iter.try_map_ast(f)?,
+            ifs: ifs.try_map_ast(f)?,
+            is_async: is_async.try_map_ast(f)?,
+        })
+    }
+}
+impl<T, U> MapAst<T, U> for ExcepthandlerKind<T> {
+    type Mapped = ExcepthandlerKind<U>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        match self {
+            ExcepthandlerKind::ExceptHandler { type_, name, body } => {
+                Ok(ExcepthandlerKind::ExceptHandler {
+                    type_: type_.try_map_ast(f)?,
+                    name: name.try_map_ast(f)?,
+                    body: body.try_map_ast(f)?,
+                })
+            }
+        }
+    }
+}
+impl<T, U> MapAst<T, U> for Arguments<T> {
+    type Mapped = Arguments<U>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        let Arguments {
+            posonlyargs,
+            args,
+            vararg,
+            kwonlyargs,
+            kw_defaults,
+            kwarg,
+            defaults,
+        } = self;
+        Ok(Arguments {
+            posonlyargs: posonlyargs.try_map_ast(f)?,
+            args: args.try_map_ast(f)?,
+            vararg: vararg.try_map_ast(f)?,
+            kwonlyargs: kwonlyargs.try_map_ast(f)?,
+            kw_defaults: kw_defaults.try_map_ast(f)?,
+            kwarg: kwarg.try_map_ast(f)?,
+            defaults: defaults.try_map_ast(f)?,
+        })
+    }
+}
+impl<T, U> MapAst<T, U> for ArgData<T> {
+    type Mapped = ArgData<U>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        let ArgData {
+            arg,
+            annotation,
+            type_comment,
+        } = self;
+        Ok(ArgData {
+            arg: arg.try_map_ast(f)?,
+            annotation: annotation.try_map_ast(f)?,
+            type_comment: type_comment.try_map_ast(f)?,
+        })
+    }
+}
+impl<T, U> MapAst<T, U> for KeywordData<T> {
+    type Mapped = KeywordData<U>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        let KeywordData { arg, value } = self;
+        Ok(KeywordData {
+            arg: arg.try_map_ast(f)?,
+            value: value.try_map_ast(f)?,
+        })
+    }
+}
+no_user!(Alias);
+impl<T, U> MapAst<T, U> for Withitem<T> {
+    type Mapped = Withitem<U>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        let Withitem {
+            context_expr,
+            optional_vars,
+        } = self;
+        Ok(Withitem {
+            context_expr: context_expr.try_map_ast(f)?,
+            optional_vars: optional_vars.try_map_ast(f)?,
+        })
+    }
+}
+no_user!(TypeIgnore);

ast/src/lib.rs

@@ -1,3 +1,7 @@
+// above for the macro
+#[macro_use]
+pub mod map_ast;
+
 mod ast_gen;
 mod constant;
 mod impls;

ast/src/map_ast.rs

@@ -0,0 +1,66 @@
+use std::convert::Infallible;
+
+pub trait MapAst<T, U>: Sized {
+    type Mapped;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E>;
+    fn map_ast<F: FnMut(T) -> U>(self, mut f: F) -> Self::Mapped {
+        let result: Result<_, Infallible> = self.try_map_ast(&mut |u| Ok(f(u)));
+        match result {
+            Ok(mapped) => mapped,
+            Err(never) => match never {},
+        }
+    }
+}
+
+macro_rules! no_user {
+    ($t:ty) => {
+        impl<T, U> MapAst<T, U> for $t {
+            type Mapped = Self;
+            #[inline]
+            fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(
+                self,
+                _f: &mut F,
+            ) -> Result<Self::Mapped, E> {
+                Ok(self)
+            }
+        }
+    };
+}
+
+no_user!(String);
+no_user!(crate::Constant);
+no_user!(crate::ConversionFlag);
+no_user!(bool);
+no_user!(usize);
+
+impl<T, U, A: MapAst<T, U>> MapAst<T, U> for crate::Located<A, T> {
+    type Mapped = crate::Located<A::Mapped, U>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        Ok(crate::Located {
+            location: self.location,
+            custom: f(self.custom)?,
+            node: self.node.try_map_ast(f)?,
+        })
+    }
+}
+
+impl<T, U, A: MapAst<T, U>> MapAst<T, U> for Vec<A> {
+    type Mapped = Vec<A::Mapped>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        self.into_iter().map(|node| node.try_map_ast(f)).collect()
+    }
+}
+
+impl<T, U, A: MapAst<T, U>> MapAst<T, U> for Option<A> {
+    type Mapped = Option<A::Mapped>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        self.map(|node| node.try_map_ast(f)).transpose()
+    }
+}
+
+impl<T, U, A: MapAst<T, U>> MapAst<T, U> for Box<A> {
+    type Mapped = Box<A::Mapped>;
+    fn try_map_ast<E, F: FnMut(T) -> Result<U, E>>(self, f: &mut F) -> Result<Self::Mapped, E> {
+        (*self).try_map_ast(f).map(Box::new)
+    }
+}

vm/src/stdlib/ast.rs

@@ -17,8 +17,8 @@ use rustpython_compiler as compile;
 use crate::builtins::{self, PyStrRef, PyTypeRef};
 use crate::function::FuncArgs;
 use crate::pyobject::{
-    BorrowValue, IdProtocol, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult,
-    PyValue, StaticType, TryFromObject, TypeProtocol,
+    BorrowValue, IdProtocol, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyResult, PyValue,
+    StaticType, TryFromObject, TypeProtocol,
 };
 use crate::vm::VirtualMachine;
 
@@ -26,13 +26,6 @@ use crate::vm::VirtualMachine;
 #[allow(clippy::all)]
 mod gen;
 
-fn node_add_location(node: &AstNodeRef, location: ast::Location, vm: &VirtualMachine) {
-    let dict = node.as_object().dict().unwrap();
-    dict.set_item("lineno", vm.ctx.new_int(location.row()), vm)
-        .unwrap();
-    dict.set_item("col_offset", vm.ctx.new_int(location.column()), vm)
-        .unwrap();
-}
 
 fn get_node_field(vm: &VirtualMachine, obj: &PyObjectRef, field: &str, typ: &str) -> PyResult {
     vm.get_attribute_opt(obj.clone(), field)?.ok_or_else(|| {
@@ -53,7 +46,6 @@ fn get_node_field_opt(
 #[pyclass(module = "_ast", name = "AST")]
 #[derive(Debug)]
 pub(crate) struct AstNode;
-type AstNodeRef = PyRef<AstNode>;
 
 #[pyimpl(flags(HAS_DICT))]
 impl AstNode {
@@ -103,6 +95,10 @@ trait Node: Sized {
     fn ast_from_object(vm: &VirtualMachine, object: PyObjectRef) -> PyResult<Self>;
 }
 
+trait NamedNode: Node {
+    const NAME: &'static str;
+}
+
 impl<T: Node> Node for Vec<T> {
     fn ast_to_object(self, vm: &VirtualMachine) -> PyObjectRef {
         vm.ctx.new_list(
@@ -144,6 +140,31 @@ impl<T: Node> Node for Option<T> {
     }
 }
 
+impl<T: NamedNode> Node for ast::Located<T> {
+    fn ast_to_object(self, vm: &VirtualMachine) -> PyObjectRef {
+        let obj = self.node.ast_to_object(vm);
+        node_add_location(&obj, self.location, vm);
+        obj
+    }
+
+    fn ast_from_object(vm: &VirtualMachine, object: PyObjectRef) -> PyResult<Self> {
+        let location = ast::Location::new(
+            Node::ast_from_object(vm, get_node_field(vm, &object, "lineno", T::NAME)?)?,
+            Node::ast_from_object(vm, get_node_field(vm, &object, "col_offset", T::NAME)?)?,
+        );
+        let node = T::ast_from_object(vm, object)?;
+        Ok(ast::Located::new(location, node))
+    }
+}
+
+fn node_add_location(node: &PyObjectRef, location: ast::Location, vm: &VirtualMachine) {
+    let dict = node.dict().unwrap();
+    dict.set_item("lineno", vm.ctx.new_int(location.row()), vm)
+        .unwrap();
+    dict.set_item("col_offset", vm.ctx.new_int(location.column()), vm)
+        .unwrap();
+}
+
 impl Node for String {
     fn ast_to_object(self, vm: &VirtualMachine) -> PyObjectRef {
         vm.ctx.new_str(self)