github.com/varialus/godfly@v0.0.0-20130904042352-1934f9f095ab/doc/articles/wiki/srcextract.go (about)

     1  // Copyright 2010 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"bytes"
     9  	"flag"
    10  	"go/ast"
    11  	"go/parser"
    12  	"go/printer"
    13  	"go/token"
    14  	"log"
    15  	"os"
    16  	"text/template"
    17  )
    18  
    19  var (
    20  	srcFn   = flag.String("src", "", "source filename")
    21  	getName = flag.String("name", "", "func/type name to output")
    22  	html    = flag.Bool("html", true, "output HTML")
    23  	showPkg = flag.Bool("pkg", false, "show package in output")
    24  )
    25  
    26  func main() {
    27  	// handle input
    28  	flag.Parse()
    29  	if *srcFn == "" || *getName == "" {
    30  		flag.Usage()
    31  		os.Exit(2)
    32  	}
    33  	// load file
    34  	fs := token.NewFileSet()
    35  	file, err := parser.ParseFile(fs, *srcFn, nil, 0)
    36  	if err != nil {
    37  		log.Fatal(err)
    38  	}
    39  	// create filter
    40  	filter := func(name string) bool {
    41  		return name == *getName
    42  	}
    43  	// filter
    44  	if !ast.FilterFile(file, filter) {
    45  		os.Exit(1)
    46  	}
    47  	// print the AST
    48  	var b bytes.Buffer
    49  	printer.Fprint(&b, fs, file)
    50  	// drop package declaration
    51  	if !*showPkg {
    52  		for {
    53  			c, err := b.ReadByte()
    54  			if c == '\n' || err != nil {
    55  				break
    56  			}
    57  		}
    58  	}
    59  	// drop leading newlines
    60  	for {
    61  		b, err := b.ReadByte()
    62  		if err != nil {
    63  			break
    64  		}
    65  		if b != '\n' {
    66  			os.Stdout.Write([]byte{b})
    67  			break
    68  		}
    69  	}
    70  	// output
    71  	if *html {
    72  		template.HTMLEscape(os.Stdout, b.Bytes())
    73  	} else {
    74  		b.WriteTo(os.Stdout)
    75  	}
    76  }