main.go

Documentation: go.uber.org/atomic/internal/gen-atomicwrapper

     1  // Copyright (c) 2020-2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  // gen-atomicwrapper generates wrapper types around other atomic types.
    22  //
    23  // It supports plugging in functions which convert the value inside the atomic
    24  // type to the user-facing value. For example,
    25  //
    26  // Given, atomic.Value and the functions,
    27  //
    28  //	func packString(string) interface{}
    29  //	func unpackString(interface{}) string
    30  //
    31  // We can run the following command:
    32  //
    33  //	gen-atomicwrapper -name String -wrapped Value \
    34  //	  -type string -pack fromString -unpack tostring
    35  //
    36  // This wil generate approximately,
    37  //
    38  //	type String struct{ v Value }
    39  //
    40  //	func (s *String) Load() string {
    41  //	  return unpackString(v.Load())
    42  //	}
    43  //
    44  //	func (s *String) Store(s string) {
    45  //	  return s.v.Store(packString(s))
    46  //	}
    47  //
    48  // The packing/unpacking logic allows the stored value to be different from
    49  // the user-facing value.
    50  package main
    51  
    52  import (
    53  	"bytes"
    54  	"embed"
    55  	"errors"
    56  	"flag"
    57  	"fmt"
    58  	"go/format"
    59  	"io"
    60  	"log"
    61  	"os"
    62  	"sort"
    63  	"strings"
    64  	"text/template"
    65  	"time"
    66  )
    67  
    68  func main() {
    69  	log.SetFlags(0)
    70  	if err := run(os.Args[1:]); err != nil {
    71  		log.Fatalf("%+v", err)
    72  	}
    73  }
    74  
    75  type stringList []string
    76  
    77  func (sl *stringList) String() string {
    78  	return strings.Join(*sl, ",")
    79  }
    80  
    81  func (sl *stringList) Set(s string) error {
    82  	for _, i := range strings.Split(s, ",") {
    83  		*sl = append(*sl, strings.TrimSpace(i))
    84  	}
    85  	return nil
    86  }
    87  
    88  func run(args []string) error {
    89  	var opts struct {
    90  		Name    string
    91  		Wrapped string
    92  		Type    string
    93  
    94  		Imports      stringList
    95  		Pack, Unpack string
    96  
    97  		CAS            bool
    98  		CompareAndSwap bool
    99  		Swap           bool
   100  		JSON           bool
   101  
   102  		File   string
   103  		ToYear int
   104  	}
   105  
   106  	opts.ToYear = time.Now().Year()
   107  
   108  	flag := flag.NewFlagSet("gen-atomicwrapper", flag.ContinueOnError)
   109  
   110  	// Required flags
   111  	flag.StringVar(&opts.Name, "name", "",
   112  		"name of the generated type (e.g. Duration)")
   113  	flag.StringVar(&opts.Wrapped, "wrapped", "",
   114  		"name of the wrapped atomic (e.g. Int64)")
   115  	flag.StringVar(&opts.Type, "type", "",
   116  		"name of the type exposed by the atomic (e.g. time.Duration)")
   117  
   118  	// Optional flags
   119  	flag.Var(&opts.Imports, "imports",
   120  		"comma separated list of imports to add")
   121  	flag.StringVar(&opts.Pack, "pack", "",
   122  		"function to transform values with before storage")
   123  	flag.StringVar(&opts.Unpack, "unpack", "",
   124  		"function to reverse packing on loading")
   125  	flag.StringVar(&opts.File, "file", "",
   126  		"output file path (default: stdout)")
   127  
   128  	// Switches for individual methods. Underlying atomics must support
   129  	// these.
   130  	flag.BoolVar(&opts.CAS, "cas", false,
   131  		"generate a deprecated `CAS(old, new) bool` method; requires -pack")
   132  	flag.BoolVar(&opts.CompareAndSwap, "compareandswap", false,
   133  		"generate a `CompareAndSwap(old, new) bool` method; requires -pack")
   134  	flag.BoolVar(&opts.Swap, "swap", false,
   135  		"generate a `Swap(new) old` method; requires -pack and -unpack")
   136  	flag.BoolVar(&opts.JSON, "json", false,
   137  		"generate `MarshalJSON/UnmarshJSON` methods")
   138  
   139  	if err := flag.Parse(args); err != nil {
   140  		return err
   141  	}
   142  
   143  	if len(opts.Name) == 0 ||
   144  		len(opts.Wrapped) == 0 ||
   145  		len(opts.Type) == 0 ||
   146  		len(opts.Pack) == 0 ||
   147  		len(opts.Unpack) == 0 {
   148  		return errors.New("flags -name, -wrapped, -pack, -unpack and -type are required")
   149  	}
   150  
   151  	if opts.CAS {
   152  		opts.CompareAndSwap = true
   153  	}
   154  
   155  	var w io.Writer = os.Stdout
   156  	if file := opts.File; len(file) > 0 {
   157  		f, err := os.Create(file)
   158  		if err != nil {
   159  			return fmt.Errorf("create %q: %v", file, err)
   160  		}
   161  		defer f.Close()
   162  
   163  		w = f
   164  	}
   165  
   166  	// Import encoding/json if needed.
   167  	if opts.JSON {
   168  		found := false
   169  		for _, imp := range opts.Imports {
   170  			if imp == "encoding/json" {
   171  				found = true
   172  				break
   173  			}
   174  		}
   175  
   176  		if !found {
   177  			opts.Imports = append(opts.Imports, "encoding/json")
   178  		}
   179  	}
   180  
   181  	sort.Strings([]string(opts.Imports))
   182  
   183  	var buff bytes.Buffer
   184  	if err := _tmpl.ExecuteTemplate(&buff, "wrapper.tmpl", opts); err != nil {
   185  		return fmt.Errorf("render template: %v", err)
   186  	}
   187  
   188  	bs, err := format.Source(buff.Bytes())
   189  	if err != nil {
   190  		return fmt.Errorf("reformat source: %v", err)
   191  	}
   192  
   193  	io.WriteString(w, "// @generated Code generated by gen-atomicwrapper.\n\n")
   194  	_, err = w.Write(bs)
   195  	return err
   196  }
   197  
   198  var (
   199  	//go:embed *.tmpl
   200  	_tmplFS embed.FS
   201  
   202  	_tmpl = template.Must(template.New("atomicwrapper").ParseFS(_tmplFS, "*.tmpl"))
   203  )
   204
View as plain text