1// Package astfmt implements `ast.Node` formatting with fmt-like API.
2package astfmt
3
4import (
5	"bytes"
6	"fmt"
7	"go/ast"
8	"go/printer"
9	"go/token"
10	"io"
11)
12
13// Println calls fmt.Println with additional support of %s format
14// for ast.Node arguments.
15//
16// Uses empty file set for AST printing.
17func Println(args ...interface{}) error {
18	return defaultPrinter.Println(args...)
19}
20
21// Fprintf calls fmt.Fprintf with additional support of %s format
22// for ast.Node arguments.
23//
24// Uses empty file set for AST printing.
25func Fprintf(w io.Writer, format string, args ...interface{}) error {
26	return defaultPrinter.Fprintf(w, format, args...)
27}
28
29// Sprintf calls fmt.Sprintf with additional support of %s format
30// for ast.Node arguments.
31//
32// Uses empty file set for AST printing.
33func Sprintf(format string, args ...interface{}) string {
34	return defaultPrinter.Sprintf(format, args...)
35}
36
37// Sprint calls fmt.Sprint with additional support of %s format
38// for ast.Node arguments.
39//
40// Uses empty file set for AST printing.
41func Sprint(args ...interface{}) string {
42	return defaultPrinter.Sprint(args...)
43}
44
45// NewPrinter returns printer that uses bound file set when printing AST nodes.
46func NewPrinter(fset *token.FileSet) *Printer {
47	return &Printer{fset: fset}
48}
49
50// Printer provides API close to fmt package for printing AST nodes.
51// Unlike freestanding functions from this package, it makes it possible
52// to associate appropriate file set for better output.
53type Printer struct {
54	fset *token.FileSet
55}
56
57// Println printer method is like Println function, but uses bound file set when printing.
58func (p *Printer) Println(args ...interface{}) error {
59	_, err := fmt.Println(wrapArgs(p.fset, args)...)
60	return err
61}
62
63// Fprintf printer method is like Fprintf function, but uses bound file set when printing.
64func (p *Printer) Fprintf(w io.Writer, format string, args ...interface{}) error {
65	_, err := fmt.Fprintf(w, format, wrapArgs(p.fset, args)...)
66	return err
67}
68
69// Sprintf printer method is like Sprintf function, but uses bound file set when printing.
70func (p *Printer) Sprintf(format string, args ...interface{}) string {
71	return fmt.Sprintf(format, wrapArgs(p.fset, args)...)
72}
73
74// Sprint printer method is like Sprint function, but uses bound file set when printing.
75func (p *Printer) Sprint(args ...interface{}) string {
76	return fmt.Sprint(wrapArgs(p.fset, args)...)
77}
78
79// defaultPrinter is used in printing functions like Println.
80// Uses empty file set.
81var defaultPrinter = NewPrinter(token.NewFileSet())
82
83// wrapArgs returns arguments slice with every ast.Node element
84// replaced with fmtNode wrapper that supports additional formatting.
85func wrapArgs(fset *token.FileSet, args []interface{}) []interface{} {
86	for i := range args {
87		if x, ok := args[i].(ast.Node); ok {
88			args[i] = fmtNode{fset: fset, node: x}
89		}
90	}
91	return args
92}
93
94type fmtNode struct {
95	fset *token.FileSet
96	node ast.Node
97}
98
99func (n fmtNode) String() string {
100	var buf bytes.Buffer
101	if err := printer.Fprint(&buf, n.fset, n.node); err != nil {
102		return fmt.Sprintf("%%!s(ast.Node=%s)", err)
103	}
104	return buf.String()
105}
106
107func (n fmtNode) GoString() string {
108	var buf bytes.Buffer
109	fmt.Fprintf(&buf, "%#v", n.node)
110	return buf.String()
111}
112