command.go

Documentation: github.com/gogo/protobuf/vanity/command

     1  // Protocol Buffers for Go with Gadgets
     2  //
     3  // Copyright (c) 2015, The GoGo Authors. All rights reserved.
     4  // http://github.com/gogo/protobuf
     5  //
     6  // Redistribution and use in source and binary forms, with or without
     7  // modification, are permitted provided that the following conditions are
     8  // met:
     9  //
    10  //     * Redistributions of source code must retain the above copyright
    11  // notice, this list of conditions and the following disclaimer.
    12  //     * Redistributions in binary form must reproduce the above
    13  // copyright notice, this list of conditions and the following disclaimer
    14  // in the documentation and/or other materials provided with the
    15  // distribution.
    16  //
    17  // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
    18  // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
    19  // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
    20  // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
    21  // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
    22  // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
    23  // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
    24  // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
    25  // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
    26  // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
    27  // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
    28  
    29  package command
    30  
    31  import (
    32  	"fmt"
    33  	"go/format"
    34  	"io/ioutil"
    35  	"os"
    36  	"strings"
    37  
    38  	_ "github.com/gogo/protobuf/plugin/compare"
    39  	_ "github.com/gogo/protobuf/plugin/defaultcheck"
    40  	_ "github.com/gogo/protobuf/plugin/description"
    41  	_ "github.com/gogo/protobuf/plugin/embedcheck"
    42  	_ "github.com/gogo/protobuf/plugin/enumstringer"
    43  	_ "github.com/gogo/protobuf/plugin/equal"
    44  	_ "github.com/gogo/protobuf/plugin/face"
    45  	_ "github.com/gogo/protobuf/plugin/gostring"
    46  	_ "github.com/gogo/protobuf/plugin/marshalto"
    47  	_ "github.com/gogo/protobuf/plugin/oneofcheck"
    48  	_ "github.com/gogo/protobuf/plugin/populate"
    49  	_ "github.com/gogo/protobuf/plugin/size"
    50  	_ "github.com/gogo/protobuf/plugin/stringer"
    51  	"github.com/gogo/protobuf/plugin/testgen"
    52  	_ "github.com/gogo/protobuf/plugin/union"
    53  	_ "github.com/gogo/protobuf/plugin/unmarshal"
    54  	"github.com/gogo/protobuf/proto"
    55  	"github.com/gogo/protobuf/protoc-gen-gogo/generator"
    56  	_ "github.com/gogo/protobuf/protoc-gen-gogo/grpc"
    57  	plugin "github.com/gogo/protobuf/protoc-gen-gogo/plugin"
    58  )
    59  
    60  func Read() *plugin.CodeGeneratorRequest {
    61  	g := generator.New()
    62  	data, err := ioutil.ReadAll(os.Stdin)
    63  	if err != nil {
    64  		g.Error(err, "reading input")
    65  	}
    66  
    67  	if err := proto.Unmarshal(data, g.Request); err != nil {
    68  		g.Error(err, "parsing input proto")
    69  	}
    70  
    71  	if len(g.Request.FileToGenerate) == 0 {
    72  		g.Fail("no files to generate")
    73  	}
    74  	return g.Request
    75  }
    76  
    77  // filenameSuffix replaces the .pb.go at the end of each filename.
    78  func GeneratePlugin(req *plugin.CodeGeneratorRequest, p generator.Plugin, filenameSuffix string) *plugin.CodeGeneratorResponse {
    79  	g := generator.New()
    80  	g.Request = req
    81  	if len(g.Request.FileToGenerate) == 0 {
    82  		g.Fail("no files to generate")
    83  	}
    84  
    85  	g.CommandLineParameters(g.Request.GetParameter())
    86  
    87  	g.WrapTypes()
    88  	g.SetPackageNames()
    89  	g.BuildTypeNameMap()
    90  	g.GeneratePlugin(p)
    91  
    92  	for i := 0; i < len(g.Response.File); i++ {
    93  		g.Response.File[i].Name = proto.String(
    94  			strings.Replace(*g.Response.File[i].Name, ".pb.go", filenameSuffix, -1),
    95  		)
    96  	}
    97  	if err := goformat(g.Response); err != nil {
    98  		g.Error(err)
    99  	}
   100  	return g.Response
   101  }
   102  
   103  func goformat(resp *plugin.CodeGeneratorResponse) error {
   104  	for i := 0; i < len(resp.File); i++ {
   105  		formatted, err := format.Source([]byte(resp.File[i].GetContent()))
   106  		if err != nil {
   107  			return fmt.Errorf("go format error: %v", err)
   108  		}
   109  		fmts := string(formatted)
   110  		resp.File[i].Content = &fmts
   111  	}
   112  	return nil
   113  }
   114  
   115  func Generate(req *plugin.CodeGeneratorRequest) *plugin.CodeGeneratorResponse {
   116  	// Begin by allocating a generator. The request and response structures are stored there
   117  	// so we can do error handling easily - the response structure contains the field to
   118  	// report failure.
   119  	g := generator.New()
   120  	g.Request = req
   121  
   122  	g.CommandLineParameters(g.Request.GetParameter())
   123  
   124  	// Create a wrapped version of the Descriptors and EnumDescriptors that
   125  	// point to the file that defines them.
   126  	g.WrapTypes()
   127  
   128  	g.SetPackageNames()
   129  	g.BuildTypeNameMap()
   130  
   131  	g.GenerateAllFiles()
   132  
   133  	if err := goformat(g.Response); err != nil {
   134  		g.Error(err)
   135  	}
   136  
   137  	testReq := proto.Clone(req).(*plugin.CodeGeneratorRequest)
   138  
   139  	testResp := GeneratePlugin(testReq, testgen.NewPlugin(), "pb_test.go")
   140  
   141  	for i := 0; i < len(testResp.File); i++ {
   142  		if strings.Contains(*testResp.File[i].Content, `//These tests are generated by github.com/gogo/protobuf/plugin/testgen`) {
   143  			g.Response.File = append(g.Response.File, testResp.File[i])
   144  		}
   145  	}
   146  
   147  	return g.Response
   148  }
   149  
   150  func Write(resp *plugin.CodeGeneratorResponse) {
   151  	g := generator.New()
   152  	// Send back the results.
   153  	data, err := proto.Marshal(resp)
   154  	if err != nil {
   155  		g.Error(err, "failed to marshal output proto")
   156  	}
   157  	_, err = os.Stdout.Write(data)
   158  	if err != nil {
   159  		g.Error(err, "failed to write output proto")
   160  	}
   161  }
   162
View as plain text