Code Generation, Pt. 3: extension code generation (#259)

jlowellwofford · web-flow · commit 7eaea1d64acd · 2021-04-19T18:04:13.000-06:00
* add code generation for extensions

* updated README to demonstrate extension building

* fixup extension generation to allow multiple extensions in one pkg
diff --git a/README.md b/README.md
@@ -154,7 +154,7 @@ $ mkdir -p modules/test
 $ cd modules/test
 <create module.yaml>
 $ kraken module generate 
-INFO[0000] module "test" generated at modules/test      
+INFO[0000] module "test" generated at "."
 ```
 
 If we selected `with_config: true`, we will need to generate the protobuf code from the provided `proto` file.  You can add some variables to `test.config.proto` first, then:
@@ -180,7 +180,27 @@ This will *only* update `test.mod.go`.  Note: you may need to make manual change
 
 ### Generating an extension
 
-Extension generation is not yet supported.
+Extensions are the least complicated to generate.  The definition file for an extension looks like:
+
+```yaml
+---
+package_url: github.com/kraken-hpc/kraken/test
+name: TestMessage
+custom_types:
+  - "MySpecialType"  
+```
+
+This will generate an extension that will be referenced as `Test.TestMessage`.  Note that we support generating multiple extensions in the same proto package using multiple definition files.  E.g., you could also have `Test.AnotherMessage` defined in another file.
+
+The procedure is similar to the others.  The default file name for extensions is `extension.yaml`:
+
+```bash
+$ mkdir -p extensions/test
+$ cd extensions/test
+<create extension.yaml>
+$ kraken extension generate 
+INFO[0000] extension "Test.TestMessage" generated at "."  
+```
 
 ## I want to get involved...
 
diff --git a/generators/extension.go b/generators/extension.go
@@ -0,0 +1,224 @@
+/* extension.go: generators for making kraken extensions
+ *
+ * Author: J. Lowell Wofford <lowell@lanl.gov>
+ *
+ * This software is open source software available under the BSD-3 license.
+ * Copyright (c) 2021, Triad National Security, LLC
+ * See LICENSE file for details.
+ */
+
+package generators
+
+import (
+	"flag"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strings"
+	"text/template"
+
+	"github.com/kraken-hpc/kraken/generators/templates"
+	"github.com/kraken-hpc/kraken/lib/util"
+	"gopkg.in/yaml.v2"
+)
+
+type ExtensionConfig struct {
+	// Global should not be specified on config; it will get overwritten regardless.
+	Global *GlobalConfigType `yaml:""`
+	// URL of package, e.g. "github.com/kraken-hpc/kraken/extensions/ipv4" (required)
+	PackageUrl string `yaml:"package_url"`
+	// Go package name (default: last element of PackageUrl)
+	PackageName string `yaml:"package_name"`
+	// Proto package name (default: camel case of PackageName)
+	ProtoPackage string `yaml:"proto_name"`
+	// Extension object name (required)
+	// More than one extension object can be declared in the same package
+	// Objects are referenced as ProtoName.Name, e.g. IPv4.IPv4OverEthernet
+	Name string `yaml:"name"`
+	// CustomTypes are an advanced feature and most people won't use them
+	// Declaring a custom type will create a stub to develop a gogo customtype on
+	// It will also include a commented example of linking a customtype in the proto
+	CustomTypes []string `yaml:"custom_types"`
+	// LowerName is intended for internal use only
+	LowerName string `yaml:""`
+}
+
+type CustomTypeConfig struct {
+	Global *GlobalConfigType `yaml:"__global,omitempty"`
+	Name   string
+}
+
+func extensionReadConfig(file string) (cfg *ExtensionConfig, err error) {
+	cfgData, err := ioutil.ReadFile(file)
+	if err != nil {
+		return nil, fmt.Errorf("could not read config file %s: %v", file, err)
+	}
+	cfg = &ExtensionConfig{}
+	if err = yaml.Unmarshal(cfgData, cfg); err != nil {
+		return nil, fmt.Errorf("failed to parse config file %s: %v", file, err)
+	}
+	// now, apply sanity checks & defaults
+	if cfg.PackageUrl == "" {
+		return nil, fmt.Errorf("package_url must be specified")
+	}
+	if cfg.Name == "" {
+		return nil, fmt.Errorf("name must be specified")
+	}
+	urlParts := util.URLToSlice(cfg.PackageUrl)
+	name := urlParts[len(urlParts)-1]
+	if cfg.PackageName == "" {
+		cfg.PackageName = name
+	}
+	if cfg.ProtoPackage == "" {
+		parts := strings.Split(cfg.PackageName, "_") // in the off chance _ is used
+		cname := ""
+		for _, s := range parts {
+			cname += strings.Title(s)
+		}
+		cfg.ProtoPackage = cname
+	}
+	cfg.LowerName = strings.ToLower(cfg.Name)
+	return
+}
+
+func extensionCompileTemplate(tplFile, outDir string, cfg *ExtensionConfig) (target string, err error) {
+	var tpl *template.Template
+	var out *os.File
+	parts := strings.Split(filepath.Base(tplFile), ".")
+	if parts[0] == "template" {
+		parts = append([]string{cfg.LowerName}, parts[1:len(parts)-1]...)
+	} else {
+		parts = parts[:len(parts)-1]
+	}
+	target = strings.Join(parts, ".")
+	dest := filepath.Join(outDir, target)
+	if _, err := os.Stat(dest); err == nil {
+		if !Global.Force {
+			return "", fmt.Errorf("refusing to overwrite file: %s (force not specified)", dest)
+		}
+	}
+	tpl = template.New(tplFile)
+	data, err := templates.Asset(tplFile)
+	if err != nil {
+		return
+	}
+	if tpl, err = tpl.Parse(string(data)); err != nil {
+		return
+	}
+	if out, err = os.Create(dest); err != nil {
+		return
+	}
+	defer out.Close()
+	err = tpl.Execute(out, cfg)
+	return
+}
+
+func extensionCompileCustomtype(outDir string, cfg *CustomTypeConfig) (target string, err error) {
+	var tpl *template.Template
+	var out *os.File
+	target = cfg.Name + ".type.go"
+	dest := filepath.Join(outDir, target)
+	if _, err := os.Stat(dest); err == nil {
+		if !Global.Force {
+			return "", fmt.Errorf("refusing to overwrite file: %s (force not specified)", dest)
+		}
+	}
+	tpl = template.New("extension/customtype.type.go.tpl")
+	data, err := templates.Asset("extension/customtype.type.go.tpl")
+	if err != nil {
+		return
+	}
+	if tpl, err = tpl.Parse(string(data)); err != nil {
+		return
+	}
+	if out, err = os.Create(dest); err != nil {
+		return
+	}
+	defer out.Close()
+	err = tpl.Execute(out, cfg)
+	return
+}
+
+func ExtensionGenerate(global *GlobalConfigType, args []string) {
+	Global = global
+	Log = global.Log
+	var configFile string
+	var outDir string
+	var help bool
+	fs := flag.NewFlagSet("extension generate", flag.ExitOnError)
+	fs.StringVar(&configFile, "c", "extension.yaml", "name of extension config file to use")
+	fs.StringVar(&outDir, "o", ".", "output directory for extension")
+	fs.BoolVar(&help, "h", false, "print this usage")
+	fs.Usage = func() {
+		fmt.Println("extension [gen]erate will generate a kraken extension based on an extension config.")
+		fmt.Println("Usage: kraken <opts> extension generate [-h] [-c <config_file>] [-o <out_dir>]")
+		fs.PrintDefaults()
+	}
+	fs.Parse(args)
+	if help {
+		fs.Usage()
+		os.Exit(0)
+	}
+	if len(fs.Args()) != 0 {
+		Log.Fatalf("unknown option: %s", fs.Args()[0])
+	}
+	stat, err := os.Stat(outDir)
+	if err == nil && !stat.IsDir() {
+		Log.Fatalf("output directory %s exists, but is not a directory", outDir)
+	}
+	if err != nil {
+		// create the dir
+		if err = os.MkdirAll(outDir, 0777); err != nil {
+			Log.Fatalf("failed to create output directory %s: %v", outDir, err)
+		}
+	}
+	cfg, err := extensionReadConfig(configFile)
+	if err != nil {
+		Log.Fatalf("failed to read config file: %v", err)
+	}
+	Log.Debugf("generating %s.%s with %d custom types",
+		cfg.ProtoPackage,
+		cfg.Name,
+		len(cfg.CustomTypes))
+	cfg.Global = global
+	// Ok, that's all the prep, now fill/write the templates
+	common := []string{
+		"extension/template.proto.tpl",
+		"extension/template.ext.go.tpl",
+		"extension/template.go.tpl",
+	}
+	for _, f := range common {
+		written, err := extensionCompileTemplate(f, outDir, cfg)
+		if err != nil {
+			Log.Fatalf("failed to write template: %v", err)
+		}
+		Log.Debugf("wrote file: %s", written)
+	}
+	if len(cfg.CustomTypes) > 0 {
+		// generate customtypes
+		ctypeDir := filepath.Join(outDir, "customtypes")
+		stat, err := os.Stat(ctypeDir)
+		if err == nil && !stat.IsDir() {
+			Log.Fatalf("customtypes directory %s exists, but is not a directory", outDir)
+		}
+		if err != nil {
+			// create the dir
+			if err = os.MkdirAll(ctypeDir, 0777); err != nil {
+				Log.Fatalf("failed to create customtypes directory %s: %v", outDir, err)
+			}
+		}
+		for _, name := range cfg.CustomTypes {
+			ctCfg := &CustomTypeConfig{
+				Global: global,
+				Name:   name,
+			}
+			written, err := extensionCompileCustomtype(ctypeDir, ctCfg)
+			if err != nil {
+				Log.Fatalf("failed to write template: %v", err)
+			}
+			Log.Debugf("wrote file: %s", written)
+		}
+	}
+	Log.Infof("extension \"%s.%s\" generated at %s", cfg.ProtoPackage, cfg.Name, outDir)
+}
diff --git a/generators/global.go b/generators/global.go
@@ -7,7 +7,7 @@
  * See LICENSE file for details.
  */
 
-//go:generate go-bindata -o templates/templates.go -fs -pkg templates -prefix templates templates/app templates/module
+//go:generate go-bindata -o templates/templates.go -fs -pkg templates -prefix templates templates/app templates/module templates/extension
 
 package generators
 
diff --git a/generators/module.go b/generators/module.go
@@ -142,8 +142,8 @@ func ModuleGenerate(global *GlobalConfigType, args []string) {
 	var outDir string
 	var help bool
 	fs := flag.NewFlagSet("module generate", flag.ExitOnError)
-	fs.StringVar(&configFile, "c", "module.yaml", "name of app config file to use")
-	fs.StringVar(&outDir, "o", ".", "output directory for app")
+	fs.StringVar(&configFile, "c", "module.yaml", "name of module config file to use")
+	fs.StringVar(&outDir, "o", ".", "output directory for module")
 	fs.BoolVar(&help, "h", false, "print this usage")
 	fs.Usage = func() {
 		fmt.Println("module [gen]erate will generate a kraken module based on a module config.")
diff --git a/generators/templates/extension/customtype.type.go.tpl b/generators/templates/extension/customtype.type.go.tpl
@@ -0,0 +1,58 @@
+// One-time generated by kraken {{ .Global.Version }}.
+// You probably want to edit this file to add values you need.
+// `kraken extension update` *will not* replace this file.
+// `kraken -f extension generate` *will* replace this file.
+
+// The generated template is supposed to act more as an example than something useable out-of-the-box
+// You'll probably need to update most of these methods to make something functional.
+
+package customtypes
+
+import (
+	"strconv"
+	"strings"
+
+	"github.com/kraken-hpc/kraken/lib/types"
+)
+
+var _ (types.ExtensionCustomType) = {{ .Name }}{}
+var _ (types.ExtensionCustomTypePtr) = (*{{ .Name }})(nil)
+
+type {{ .Name }} struct {
+	// TODO: you might want this type to be based on some other kind of data
+	//       if you change this, all of the methods will need to change appropriately
+	s string
+}
+
+// Marshal controls how the type is converted to binary
+func (t {{ .Name }}) Marshal() ([]byte, error) {
+	return []byte(t.s), nil
+}
+
+// MarshalTo is like Marshal, but writes the result to `data`
+func (t *{{ .Name }}) MarshalTo(data []byte) (int, error) {
+	copy(data, []byte(t.s))
+	return len(t.s), nil
+}
+
+// Unmarshal converts an encoded []byte to populate the type
+func (t *{{ .Name }}) Unmarshal(data []byte) error {
+	t.s = string(data)
+	return nil
+}
+
+// Size returns the size (encoded) of the type
+func (t *{{ .Name }}) Size() int {
+	return len(t.s)
+}
+
+// MarshalJSON writes the type to JSON format in the form of a byte string
+func (t {{ .Name }}) MarshalJSON() (j []byte, e error) {
+	return []byte(strconv.Quote(t.s)), nil
+}
+
+// UnmarshalJSON takes byte string JSON to populate the type
+func (t *{{ .Name }}) UnmarshalJSON(data []byte) error {
+	t.s = strings.Trim(string(data), "\"'")
+	return nil
+}
diff --git a/generators/templates/extension/template.ext.go.tpl b/generators/templates/extension/template.ext.go.tpl
@@ -0,0 +1,20 @@
+// Generated by kraken {{- .Global.Version }}. DO NOT EDIT
+// To update, run `kraken extension update`
+package {{ .PackageName }}
+
+import (
+	"github.com/kraken-hpc/kraken/core"
+)
+
+//////////////////////////
+// Boilerplate Methods //
+////////////////////////
+
+
+func (*{{ .Name }}) Name() string {
+	return "type.googleapis.com/{{ .ProtoPackage }}.{{ .Name }}" 
+}
+
+func init() {
+	core.Registry.RegisterExtension(&{{ .Name }}{})
+}
diff --git a/generators/templates/extension/template.go.tpl b/generators/templates/extension/template.go.tpl
@@ -0,0 +1,28 @@
+// One-time generated by kraken {{ .Global.Version }}.
+// You probably want to edit this file to add values you need.
+// `kraken extension update` *will not* replace this file.
+// `kraken -f extension generate` *will* replace this file.
+// To generate your config protobuf run `go generate` (requires `protoc` and `gogo-protobuf`)
+
+//go:generate protoc -I $GOPATH/src -I . --gogo_out=plugins=grpc:. {{ .LowerName }}.proto
+
+package {{ .PackageName }}
+
+import (
+	"github.com/kraken-hpc/kraken/lib/types"
+)
+
+/////////////////////////
+// {{ .Name }} Object //
+///////////////////////
+
+// This null declaration ensures that we adhere to the interface
+var _ types.Extension = (*{{ .Name }})(nil)
+
+// New creates a new initialized instance of an extension
+func (i *{{ .Name }}) New() (m types.Message) {
+	m = &{{ .Name }}{}
+	// TODO: you can add extra actions that should happen when a new instance is created here
+	//       for instance, you could set default values.
+	return
+}
diff --git a/generators/templates/extension/template.proto.tpl b/generators/templates/extension/template.proto.tpl
diff --git a/generators/templates/templates.go b/generators/templates/templates.go
diff --git a/kraken.go b/kraken.go
