├── .travis.yml ├── xast_example_test.go ├── README.md └── xast.go /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | sudo: false 3 | 4 | go: 5 | - 1.8 6 | - 1.9 7 | - tip 8 | 9 | script: 10 | - go test -v ./... 11 | -------------------------------------------------------------------------------- /xast_example_test.go: -------------------------------------------------------------------------------- 1 | package xast_test 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "go/ast" 7 | "go/parser" 8 | "go/printer" 9 | "go/token" 10 | "strings" 11 | 12 | "github.com/OneOfOne/xast" 13 | ) 14 | 15 | func ExampleWalk() { 16 | src := ` 17 | package main 18 | 19 | // Foo is a foo! 20 | type Foo struct{} 21 | 22 | // NotFoo is not a foo! 23 | type NotFoo struct{} 24 | 25 | // DeleteMe needs to be deleted with this comment. 26 | func DeleteMe() {} 27 | 28 | // DeleteMeToo says hi. 29 | func DeleteMeToo() {} 30 | 31 | // GoodBoy is a good boy, yes he is! 32 | func GoodBoy() { 33 | var nf NotFoo 34 | _ = nf 35 | } 36 | ` 37 | 38 | fset := token.NewFileSet() 39 | file, err := parser.ParseFile(fset, "foo.go", src, parser.ParseComments) 40 | if err != nil { 41 | panic(err) 42 | } 43 | 44 | rewriteFn := func(n *xast.Node) *xast.Node { 45 | switch x := n.Node().(type) { 46 | case *ast.TypeSpec: 47 | if x.Name.Name == "Foo" { 48 | x.Name.Name = "Bar" 49 | // remove Foo's comment. 50 | n.Parent().Node().(*ast.GenDecl).Doc.List = nil 51 | } 52 | case *ast.CommentGroup: 53 | if strings.Contains(x.Text(), "NotFoo") { 54 | x.List[0].Text = "// NotFoo got pwned." 55 | } 56 | return n.Break() // won't delete the node but Walk won't go down its children list. 57 | case *ast.FuncDecl: 58 | switch x.Name.Name { 59 | case "DeleteMe", "DeleteMeToo": 60 | return n.Delete() 61 | case "GoodBoy": 62 | x.Doc.List = nil // remove the goodboy's comment :-/ 63 | } 64 | } 65 | 66 | return n 67 | } 68 | 69 | var buf bytes.Buffer 70 | printer.Fprint(&buf, fset, xast.Walk(file, rewriteFn)) 71 | fmt.Println(buf.String()) 72 | // Output: 73 | // package main 74 | // 75 | // type Bar struct{} 76 | // 77 | // // NotFoo got pwned. 78 | // type NotFoo struct{} 79 | // 80 | // func GoodBoy() { 81 | // var nf NotFoo 82 | // _ = nf 83 | // } 84 | } 85 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # xast [![GoDoc](http://godoc.org/github.com/OneOfOne/genx?status.svg)](http://godoc.org/github.com/OneOfOne/xast) [![Build Status](https://travis-ci.org/OneOfOne/genx.svg?branch=master)](https://travis-ci.org/OneOfOne/xast) 2 | 3 | xast provides `Walk()/WalkNode()` functions, similar to [`astrewrite.Walk`](https://godoc.org/github.com/fatih/astrewrite#example-Walk) from the 4 | [astrewrite](https://github.com/fatih/astrewrite) package. The main difference is that the passed walk function can also 5 | check a node's parent. 6 | 7 | Also deleting a node will automatically remove any assocociated comments with it. 8 | 9 | Note that this is similar to https://go-review.googlesource.com/c/go/+/55790 but the API itself is a lot simpler. 10 | 11 | # Example 12 | 13 | ```go 14 | package main 15 | 16 | import ( 17 | "bytes" 18 | "fmt" 19 | "go/ast" 20 | "go/parser" 21 | "go/printer" 22 | "go/token" 23 | 24 | "github.com/OneOfOne/xast" 25 | ) 26 | 27 | func main() { 28 | src := `package main 29 | 30 | // Foo is a foo! 31 | type Foo struct{} 32 | 33 | // NotFoo is not a foo! 34 | type NotFoo struct{} 35 | ` 36 | 37 | fset := token.NewFileSet() 38 | file, err := parser.ParseFile(fset, "foo.go", src, parser.ParseComments) 39 | if err != nil { 40 | panic(err) 41 | } 42 | 43 | rewriteFn := func(n *xast.Node) *xast.Node { 44 | switch x := n.Node().(type) { 45 | case *ast.TypeSpec: 46 | if x.Name.Name == "Foo" { 47 | x.Name.Name = "Bar" 48 | } 49 | case *ast.CommentGroup: 50 | // check n.Text()? 51 | x.List = nil 52 | return n.Break() // don't delete the node but don't go down it's children list. 53 | 54 | // or if you want to remove a single comment out of a group 55 | case *ast.Comment: // won't ever get here since we return n.Break() from case *ast.CommentGroup. 56 | panic("can't get here now") 57 | return n.Nil() // delete this node 58 | 59 | } 60 | 61 | return n 62 | } 63 | 64 | var buf bytes.Buffer 65 | printer.Fprint(&buf, fset, xast.Walk(file, rewriteFn)) 66 | fmt.Println(buf.String()) 67 | } 68 | ``` 69 | 70 | -------------------------------------------------------------------------------- /xast.go: -------------------------------------------------------------------------------- 1 | package xast 2 | 3 | import ( 4 | "fmt" 5 | "go/ast" 6 | "reflect" 7 | ) 8 | 9 | // NewNode returns a new node with the given parent and ast.Node. 10 | func NewNode(parent *Node, cur ast.Node) *Node { 11 | return &Node{p: parent, n: cur} 12 | } 13 | 14 | // Node holds the current ast.Node and a parent *Node. 15 | type Node struct { 16 | p *Node 17 | n ast.Node 18 | delete bool 19 | skip bool 20 | } 21 | 22 | // Parent returns Parent Node. 23 | func (n *Node) Parent() *Node { 24 | if n == nil { 25 | return nil 26 | } 27 | return n.p 28 | } 29 | 30 | // Node returns the current ast.Node. 31 | func (n *Node) Node() ast.Node { 32 | if n == nil { 33 | return nil 34 | } 35 | return n.n 36 | } 37 | 38 | // SetNode replaces the current ast.Node. 39 | func (n *Node) SetNode(nn ast.Node) *Node { 40 | if n != nil { 41 | n.n = nn 42 | } 43 | return n 44 | } 45 | 46 | // Delete marks the node as nil and returns it, making Walk delete it from its parent. 47 | func (n *Node) Delete() *Node { 48 | if n != nil { 49 | n.delete = true 50 | } 51 | return n 52 | } 53 | 54 | // Break skips the current node from farther processing. 55 | func (n *Node) Break() *Node { 56 | if n != nil { 57 | n.skip = true 58 | } 59 | return n 60 | } 61 | 62 | // Canceled returns true if n is nil or Break/Delete got called. 63 | func (n *Node) Canceled() bool { 64 | return n == nil || n.skip || n.delete || n.n == nil 65 | } 66 | 67 | func (n *Node) assign(dst interface{}) (assigned bool) { 68 | rv := reflect.ValueOf(dst).Elem() 69 | if assigned = n != nil && !n.delete && n.n != nil; assigned { 70 | rv.Set(reflect.ValueOf(n.n)) 71 | } else { 72 | rv.Set(reflect.Zero(rv.Type())) 73 | } 74 | return 75 | } 76 | 77 | // WalkFunc describes a function to be called for each node during a Walk. 78 | // The returned node can be used to rewrite the AST. 79 | type WalkFunc func(*Node) *Node 80 | 81 | func isNil(v interface{}) bool { 82 | rv := reflect.ValueOf(v) 83 | return !rv.IsValid() || rv.IsNil() 84 | } 85 | 86 | // Walk calls WalkNode() with the given node and returns the result. 87 | func Walk(root ast.Node, fn WalkFunc) ast.Node { 88 | return WalkNode(&Node{n: root}, fn).Node() 89 | } 90 | 91 | // WalkNode traverses an AST in depth-first order. 92 | // Panics if you call node.SetNode(new-node) the returned type is not the same type as the original one. 93 | func WalkNode(node *Node, fn WalkFunc) *Node { 94 | if isNil(node.Node()) { 95 | return node 96 | } 97 | 98 | if node = fn(node); node.Canceled() { 99 | return node 100 | } 101 | 102 | // walk children 103 | // (the order of the cases matches the order 104 | // of the corresponding node types in ast.go) 105 | switch n := node.n.(type) { 106 | case *ast.CommentGroup: 107 | out := n.List[:0] 108 | for _, c := range n.List { 109 | if WalkNode(&Node{p: node, n: c}, fn).assign(&c) { 110 | out = append(out, c) 111 | } 112 | } 113 | n.List = out 114 | 115 | case *ast.Field: 116 | n.Names = walkIdentList(node, n.Names, fn) 117 | if !WalkNode(&Node{p: node, n: n.Type}, fn).assign(&n.Type) { 118 | return node.Delete() 119 | } 120 | 121 | WalkNode(&Node{p: node, n: n.Tag}, fn).assign(&n.Tag) 122 | WalkNode(&Node{p: node, n: n.Doc}, fn).assign(&n.Doc) 123 | WalkNode(&Node{p: node, n: n.Comment}, fn).assign(&n.Comment) 124 | 125 | case *ast.FieldList: 126 | if len(n.List) == 0 { 127 | break 128 | } 129 | out := n.List[:0] 130 | for i, f := range n.List { 131 | if WalkNode(&Node{p: node, n: f}, fn).assign(&f) { 132 | out = append(out, f) 133 | } else { 134 | nukeComments(n.List[i]) 135 | } 136 | } 137 | if n.List = out; len(n.List) == 0 { 138 | return node.Delete() 139 | } 140 | 141 | case *ast.Ellipsis: 142 | if !WalkNode(&Node{p: node, n: n.Elt}, fn).assign(&n.Elt) { 143 | return node.Delete() 144 | } 145 | 146 | case *ast.FuncLit: 147 | if !WalkNode(&Node{p: node, n: n.Type}, fn).assign(&n.Type) { 148 | return node.Delete() 149 | } 150 | 151 | WalkNode(&Node{p: node, n: n.Body}, fn).assign(&n.Body) 152 | 153 | case *ast.CompositeLit: 154 | WalkNode(&Node{p: node, n: n.Type}, fn).assign(&n.Type) 155 | n.Elts = walkExprList(node, n.Elts, fn) 156 | 157 | case *ast.ParenExpr: 158 | WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) 159 | 160 | case *ast.SelectorExpr: 161 | WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) 162 | WalkNode(&Node{p: node, n: n.Sel}, fn).assign(&n.Sel) 163 | 164 | case *ast.IndexExpr: 165 | WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) 166 | WalkNode(&Node{p: node, n: n.Index}, fn).assign(&n.Index) 167 | 168 | case *ast.SliceExpr: 169 | WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) 170 | WalkNode(&Node{p: node, n: n.Low}, fn).assign(&n.Low) 171 | WalkNode(&Node{p: node, n: n.High}, fn).assign(&n.High) 172 | WalkNode(&Node{p: node, n: n.Max}, fn).assign(&n.Max) 173 | 174 | case *ast.TypeAssertExpr: 175 | WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) 176 | WalkNode(&Node{p: node, n: n.Type}, fn).assign(&n.Type) 177 | 178 | case *ast.CallExpr: 179 | if !WalkNode(&Node{p: node, n: n.Fun}, fn).assign(&n.Fun) { 180 | return node.Delete() 181 | } 182 | n.Args = walkExprList(node, n.Args, fn) 183 | 184 | case *ast.StarExpr: 185 | WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) 186 | 187 | case *ast.UnaryExpr: 188 | WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) 189 | 190 | case *ast.BinaryExpr: 191 | WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) 192 | WalkNode(&Node{p: node, n: n.Y}, fn).assign(&n.Y) 193 | 194 | case *ast.KeyValueExpr: 195 | WalkNode(&Node{p: node, n: n.Key}, fn).assign(&n.Key) 196 | WalkNode(&Node{p: node, n: n.Value}, fn).assign(&n.Value) 197 | 198 | case *ast.ArrayType: 199 | WalkNode(&Node{p: node, n: n.Len}, fn).assign(&n.Len) 200 | if !WalkNode(&Node{p: node, n: n.Elt}, fn).assign(&n.Elt) { 201 | return node.Delete() 202 | } 203 | 204 | case *ast.StructType: 205 | if !WalkNode(&Node{p: node, n: n.Fields}, fn).assign(&n.Fields) { 206 | return node.Delete() 207 | } 208 | 209 | case *ast.FuncType: 210 | WalkNode(&Node{p: node, n: n.Params}, fn).assign(&n.Params) 211 | WalkNode(&Node{p: node, n: n.Results}, fn).assign(&n.Results) 212 | 213 | case *ast.InterfaceType: 214 | WalkNode(&Node{p: node, n: n.Methods}, fn).assign(&n.Methods) 215 | 216 | case *ast.MapType: 217 | if !WalkNode(&Node{p: node, n: n.Key}, fn).assign(&n.Key) { 218 | return node.Delete() 219 | } 220 | if !WalkNode(&Node{p: node, n: n.Value}, fn).assign(&n.Value) { 221 | return node.Delete() 222 | } 223 | 224 | case *ast.ChanType: 225 | if !WalkNode(&Node{p: node, n: n.Value}, fn).assign(&n.Value) { 226 | return node.Delete() 227 | } 228 | 229 | case *ast.DeclStmt: 230 | if !WalkNode(&Node{p: node, n: n.Decl}, fn).assign(&n.Decl) { 231 | return node.Delete() 232 | } 233 | 234 | case *ast.LabeledStmt: 235 | WalkNode(&Node{p: node, n: n.Label}, fn).assign(&n.Label) 236 | WalkNode(&Node{p: node, n: n.Stmt}, fn).assign(&n.Stmt) 237 | 238 | case *ast.ExprStmt: 239 | if !WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) { 240 | return node.Delete() 241 | } 242 | 243 | case *ast.SendStmt: 244 | WalkNode(&Node{p: node, n: n.Chan}, fn).assign(&n.Chan) 245 | WalkNode(&Node{p: node, n: n.Value}, fn).assign(&n.Value) 246 | 247 | case *ast.IncDecStmt: 248 | WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) 249 | 250 | case *ast.AssignStmt: 251 | n.Lhs = walkExprList(node, n.Lhs, fn) 252 | n.Rhs = walkExprList(node, n.Rhs, fn) 253 | 254 | case *ast.GoStmt: 255 | WalkNode(&Node{p: node, n: n.Call}, fn).assign(&n.Call) 256 | 257 | case *ast.DeferStmt: 258 | WalkNode(&Node{p: node, n: n.Call}, fn).assign(&n.Call) 259 | 260 | case *ast.ReturnStmt: 261 | n.Results = walkExprList(node, n.Results, fn) 262 | 263 | case *ast.BranchStmt: 264 | 265 | WalkNode(&Node{p: node, n: n.Label}, fn).assign(&n.Label) 266 | 267 | case *ast.BlockStmt: 268 | n.List = walkStmtList(node, n.List, fn) 269 | 270 | case *ast.IfStmt: 271 | WalkNode(&Node{p: node, n: n.Init}, fn).assign(&n.Init) 272 | WalkNode(&Node{p: node, n: n.Cond}, fn).assign(&n.Cond) 273 | WalkNode(&Node{p: node, n: n.Body}, fn).assign(&n.Body) 274 | WalkNode(&Node{p: node, n: n.Else}, fn).assign(&n.Else) 275 | 276 | case *ast.CaseClause: 277 | n.List = walkExprList(node, n.List, fn) 278 | n.Body = walkStmtList(node, n.Body, fn) 279 | 280 | case *ast.SwitchStmt: 281 | 282 | WalkNode(&Node{p: node, n: n.Init}, fn).assign(&n.Init) 283 | WalkNode(&Node{p: node, n: n.Tag}, fn).assign(&n.Tag) 284 | WalkNode(&Node{p: node, n: n.Body}, fn).assign(&n.Body) 285 | 286 | case *ast.TypeSwitchStmt: 287 | WalkNode(&Node{p: node, n: n.Init}, fn).assign(&n.Init) 288 | WalkNode(&Node{p: node, n: n.Assign}, fn).assign(&n.Assign) 289 | WalkNode(&Node{p: node, n: n.Body}, fn).assign(&n.Body) 290 | 291 | case *ast.CommClause: 292 | WalkNode(&Node{p: node, n: n.Comm}, fn).assign(&n.Comm) 293 | n.Body = walkStmtList(node, n.Body, fn) 294 | 295 | case *ast.SelectStmt: 296 | WalkNode(&Node{p: node, n: n.Body}, fn).assign(&n.Body) 297 | 298 | case *ast.ForStmt: 299 | WalkNode(&Node{p: node, n: n.Init}, fn).assign(&n.Init) 300 | WalkNode(&Node{p: node, n: n.Cond}, fn).assign(&n.Cond) 301 | WalkNode(&Node{p: node, n: n.Post}, fn).assign(&n.Post) 302 | WalkNode(&Node{p: node, n: n.Body}, fn).assign(&n.Body) 303 | 304 | case *ast.RangeStmt: 305 | WalkNode(&Node{p: node, n: n.Key}, fn).assign(&n.Key) 306 | WalkNode(&Node{p: node, n: n.Value}, fn).assign(&n.Value) 307 | WalkNode(&Node{p: node, n: n.X}, fn).assign(&n.X) 308 | WalkNode(&Node{p: node, n: n.Body}, fn).assign(&n.Body) 309 | 310 | case *ast.ImportSpec: 311 | WalkNode(&Node{p: node, n: n.Doc}, fn).assign(&n.Doc) 312 | WalkNode(&Node{p: node, n: n.Name}, fn).assign(&n.Name) 313 | WalkNode(&Node{p: node, n: n.Path}, fn).assign(&n.Path) 314 | WalkNode(&Node{p: node, n: n.Comment}, fn).assign(&n.Comment) 315 | 316 | case *ast.ValueSpec: 317 | n.Names = walkIdentList(node, n.Names, fn) 318 | WalkNode(&Node{p: node, n: n.Type}, fn).assign(&n.Type) 319 | n.Values = walkExprList(node, n.Values, fn) 320 | WalkNode(&Node{p: node, n: n.Doc}, fn).assign(&n.Doc) 321 | WalkNode(&Node{p: node, n: n.Comment}, fn).assign(&n.Comment) 322 | 323 | case *ast.TypeSpec: 324 | WalkNode(&Node{p: node, n: n.Name}, fn) 325 | WalkNode(&Node{p: node, n: n.Type}, fn) 326 | WalkNode(&Node{p: node, n: n.Comment}, fn).assign(&n.Comment) 327 | 328 | case *ast.GenDecl: 329 | if n.Specs = walkSpecList(node, n.Specs, fn); len(n.Specs) == 0 { 330 | return node.Delete() 331 | } 332 | 333 | WalkNode(&Node{p: node, n: n.Doc}, fn).assign(&n.Doc) 334 | case *ast.FuncDecl: 335 | if !WalkNode(&Node{p: node, n: n.Recv}, fn).assign(&n.Recv) { 336 | return node.Delete() 337 | } 338 | WalkNode(&Node{p: node, n: n.Name}, fn).assign(&n.Name) 339 | WalkNode(&Node{p: node, n: n.Type}, fn).assign(&n.Type) 340 | WalkNode(&Node{p: node, n: n.Body}, fn).assign(&n.Body) 341 | WalkNode(&Node{p: node, n: n.Doc}, fn).assign(&n.Doc) 342 | 343 | case *ast.File: 344 | WalkNode(&Node{p: node, n: n.Doc}, fn).assign(&n.Doc) 345 | WalkNode(&Node{p: node, n: n.Name}, fn).assign(&n.Name) 346 | n.Decls = walkDeclList(node, n.Decls, fn) 347 | 348 | // don't walk n.Comments - they have been 349 | // visited already through the individual 350 | // nodes 351 | 352 | case *ast.Package: 353 | for _, f := range n.Files { 354 | WalkNode(&Node{p: node, n: f}, fn).assign(&f) 355 | } 356 | 357 | case *ast.BadStmt, *ast.BadDecl, *ast.BadExpr, *ast.Ident, 358 | *ast.BasicLit, *ast.Comment, *ast.EmptyStmt: 359 | // nothing to do 360 | 361 | default: 362 | panic(fmt.Sprintf("ast.Walk: unexpected node type %T", n)) 363 | } 364 | 365 | return node 366 | } 367 | 368 | func nukeComments(root ast.Node) { 369 | if root == nil { 370 | return 371 | } 372 | ast.Inspect(root, func(n ast.Node) bool { 373 | if isNil(n) { 374 | return false 375 | } 376 | 377 | switch n := n.(type) { 378 | case *ast.CommentGroup: 379 | n.List = nil 380 | return false 381 | case nil: 382 | return false 383 | default: 384 | return true 385 | } 386 | }) 387 | } 388 | 389 | func walkIdentList(node *Node, list []*ast.Ident, fn WalkFunc) (out []*ast.Ident) { 390 | out = list[:0] 391 | for i, x := range list { 392 | if WalkNode(&Node{p: node, n: x}, fn).assign(&x) { 393 | out = append(out, x) 394 | } else { 395 | nukeComments(list[i]) 396 | } 397 | } 398 | return 399 | } 400 | 401 | func walkExprList(node *Node, list []ast.Expr, fn WalkFunc) (out []ast.Expr) { 402 | out = list[:0] 403 | for i, x := range list { 404 | if WalkNode(&Node{p: node, n: x}, fn).assign(&x) { 405 | out = append(out, x) 406 | } else { 407 | nukeComments(list[i]) 408 | } 409 | } 410 | return 411 | } 412 | 413 | func walkStmtList(node *Node, list []ast.Stmt, fn WalkFunc) (out []ast.Stmt) { 414 | out = list[:0] 415 | for i, x := range list { 416 | if WalkNode(&Node{p: node, n: x}, fn).assign(&x) { 417 | out = append(out, x) 418 | } else { 419 | nukeComments(list[i]) 420 | } 421 | } 422 | return 423 | } 424 | 425 | func walkDeclList(node *Node, list []ast.Decl, fn WalkFunc) (out []ast.Decl) { 426 | out = list[:0] 427 | for i, x := range list { 428 | if WalkNode(&Node{p: node, n: x}, fn).assign(&x) { 429 | out = append(out, x) 430 | } else { 431 | nukeComments(list[i]) 432 | } 433 | } 434 | return 435 | } 436 | 437 | func walkSpecList(node *Node, list []ast.Spec, fn WalkFunc) (out []ast.Spec) { 438 | out = list[:0] 439 | for i, x := range list { 440 | if WalkNode(&Node{p: node, n: x}, fn).assign(&x) { 441 | out = append(out, x) 442 | } else { 443 | nukeComments(list[i]) 444 | } 445 | } 446 | 447 | return 448 | } 449 | --------------------------------------------------------------------------------