1/*
2Copyright 2015 The Kubernetes Authors.
3
4Licensed under the Apache License, Version 2.0 (the "License");
5you may not use this file except in compliance with the License.
6You may obtain a copy of the License at
7
8    http://www.apache.org/licenses/LICENSE-2.0
9
10Unless required by applicable law or agreed to in writing, software
11distributed under the License is distributed on an "AS IS" BASIS,
12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13See the License for the specific language governing permissions and
14limitations under the License.
15*/
16
17// Package args has common command-line flags for generation programs.
18package args
19
20import (
21	"bytes"
22	goflag "flag"
23	"fmt"
24	"io/ioutil"
25	"os"
26	"path"
27	"path/filepath"
28	"strconv"
29	"strings"
30	"time"
31
32	"k8s.io/gengo/generator"
33	"k8s.io/gengo/namer"
34	"k8s.io/gengo/parser"
35	"k8s.io/gengo/types"
36
37	"github.com/spf13/pflag"
38)
39
40// Default returns a defaulted GeneratorArgs. You may change the defaults
41// before calling AddFlags.
42func Default() *GeneratorArgs {
43	return &GeneratorArgs{
44		OutputBase:                 DefaultSourceTree(),
45		GoHeaderFilePath:           filepath.Join(DefaultSourceTree(), "k8s.io/gengo/boilerplate/boilerplate.go.txt"),
46		GeneratedBuildTag:          "ignore_autogenerated",
47		GeneratedByCommentTemplate: "// Code generated by GENERATOR_NAME. DO NOT EDIT.",
48		defaultCommandLineFlags:    true,
49	}
50}
51
52// GeneratorArgs has arguments that are passed to generators.
53type GeneratorArgs struct {
54	// Which directories to parse.
55	InputDirs []string
56
57	// Source tree to write results to.
58	OutputBase string
59
60	// Package path within the source tree.
61	OutputPackagePath string
62
63	// Output file name.
64	OutputFileBaseName string
65
66	// Where to get copyright header text.
67	GoHeaderFilePath string
68
69	// If GeneratedByCommentTemplate is set, generate a "Code generated by" comment
70	// below the bloilerplate, of the format defined by this string.
71	// Any instances of "GENERATOR_NAME" will be replaced with the name of the code generator.
72	GeneratedByCommentTemplate string
73
74	// If true, only verify, don't write anything.
75	VerifyOnly bool
76
77	// If true, include *_test.go files
78	IncludeTestFiles bool
79
80	// GeneratedBuildTag is the tag used to identify code generated by execution
81	// of this type. Each generator should use a different tag, and different
82	// groups of generators (external API that depends on Kube generations) should
83	// keep tags distinct as well.
84	GeneratedBuildTag string
85
86	// Any custom arguments go here
87	CustomArgs interface{}
88
89	// Whether to use default command line flags
90	defaultCommandLineFlags bool
91}
92
93// WithoutDefaultFlagParsing disables implicit addition of command line flags and parsing.
94func (g *GeneratorArgs) WithoutDefaultFlagParsing() *GeneratorArgs {
95	g.defaultCommandLineFlags = false
96	return g
97}
98
99func (g *GeneratorArgs) AddFlags(fs *pflag.FlagSet) {
100	fs.StringSliceVarP(&g.InputDirs, "input-dirs", "i", g.InputDirs, "Comma-separated list of import paths to get input types from.")
101	fs.StringVarP(&g.OutputBase, "output-base", "o", g.OutputBase, "Output base; defaults to $GOPATH/src/ or ./ if $GOPATH is not set.")
102	fs.StringVarP(&g.OutputPackagePath, "output-package", "p", g.OutputPackagePath, "Base package path.")
103	fs.StringVarP(&g.OutputFileBaseName, "output-file-base", "O", g.OutputFileBaseName, "Base name (without .go suffix) for output files.")
104	fs.StringVarP(&g.GoHeaderFilePath, "go-header-file", "h", g.GoHeaderFilePath, "File containing boilerplate header text. The string YEAR will be replaced with the current 4-digit year.")
105	fs.BoolVar(&g.VerifyOnly, "verify-only", g.VerifyOnly, "If true, only verify existing output, do not write anything.")
106	fs.StringVar(&g.GeneratedBuildTag, "build-tag", g.GeneratedBuildTag, "A Go build tag to use to identify files generated by this command. Should be unique.")
107}
108
109// LoadGoBoilerplate loads the boilerplate file passed to --go-header-file.
110func (g *GeneratorArgs) LoadGoBoilerplate() ([]byte, error) {
111	b, err := ioutil.ReadFile(g.GoHeaderFilePath)
112	if err != nil {
113		return nil, err
114	}
115	b = bytes.Replace(b, []byte("YEAR"), []byte(strconv.Itoa(time.Now().Year())), -1)
116
117	if g.GeneratedByCommentTemplate != "" {
118		if len(b) != 0 {
119			b = append(b, byte('\n'))
120		}
121		generatorName := path.Base(os.Args[0])
122		generatedByComment := strings.Replace(g.GeneratedByCommentTemplate, "GENERATOR_NAME", generatorName, -1)
123		s := fmt.Sprintf("%s\n\n", generatedByComment)
124		b = append(b, []byte(s)...)
125	}
126	return b, nil
127}
128
129// NewBuilder makes a new parser.Builder and populates it with the input
130// directories.
131func (g *GeneratorArgs) NewBuilder() (*parser.Builder, error) {
132	b := parser.New()
133
134	// flag for including *_test.go
135	b.IncludeTestFiles = g.IncludeTestFiles
136
137	// Ignore all auto-generated files.
138	b.AddBuildTags(g.GeneratedBuildTag)
139
140	for _, d := range g.InputDirs {
141		var err error
142		if strings.HasSuffix(d, "/...") {
143			err = b.AddDirRecursive(strings.TrimSuffix(d, "/..."))
144		} else {
145			err = b.AddDir(d)
146		}
147		if err != nil {
148			return nil, fmt.Errorf("unable to add directory %q: %v", d, err)
149		}
150	}
151	return b, nil
152}
153
154// InputIncludes returns true if the given package is a (sub) package of one of
155// the InputDirs.
156func (g *GeneratorArgs) InputIncludes(p *types.Package) bool {
157	for _, dir := range g.InputDirs {
158		d := dir
159		if strings.HasSuffix(d, "...") {
160			d = strings.TrimSuffix(d, "...")
161		}
162		if strings.HasPrefix(p.Path, d) {
163			return true
164		}
165	}
166	return false
167}
168
169// DefaultSourceTree returns the /src directory of the first entry in $GOPATH.
170// If $GOPATH is empty, it returns "./". Useful as a default output location.
171func DefaultSourceTree() string {
172	paths := strings.Split(os.Getenv("GOPATH"), string(filepath.ListSeparator))
173	if len(paths) > 0 && len(paths[0]) > 0 {
174		return filepath.Join(paths[0], "src")
175	}
176	return "./"
177}
178
179// Execute implements main().
180// If you don't need any non-default behavior, use as:
181// args.Default().Execute(...)
182func (g *GeneratorArgs) Execute(nameSystems namer.NameSystems, defaultSystem string, pkgs func(*generator.Context, *GeneratorArgs) generator.Packages) error {
183	if g.defaultCommandLineFlags {
184		g.AddFlags(pflag.CommandLine)
185		pflag.CommandLine.AddGoFlagSet(goflag.CommandLine)
186		pflag.Parse()
187	}
188
189	b, err := g.NewBuilder()
190	if err != nil {
191		return fmt.Errorf("Failed making a parser: %v", err)
192	}
193
194	// pass through the flag on whether to include *_test.go files
195	b.IncludeTestFiles = g.IncludeTestFiles
196
197	c, err := generator.NewContext(b, nameSystems, defaultSystem)
198	if err != nil {
199		return fmt.Errorf("Failed making a context: %v", err)
200	}
201
202	c.Verify = g.VerifyOnly
203	packages := pkgs(c, g)
204	if err := c.ExecutePackages(g.OutputBase, packages); err != nil {
205		return fmt.Errorf("Failed executing generator: %v", err)
206	}
207
208	return nil
209}
210