generator.go

package multiswagger

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"gopkg.in/yaml.v2"

	"github.com/goadesign/goa/design"
	"github.com/goadesign/goa/goagen/codegen"
	"github.com/goadesign/goa/goagen/gen_swagger"
	"github.com/goadesign/goa/goagen/utils"
)

// Generator is the swagger code generator.
type Generator struct {
	API      *design.APIDefinition // The API definition
	OutDir   string                // Path to output directory
	genfiles []string              // Generated files
}

// Generate is the generator entry point called by the meta generator.
func Generate() (files []string, err error) {
	var outDir, ver string
	set := flag.NewFlagSet("swagger", flag.PanicOnError)
	set.StringVar(&outDir, "out", "", "")
	set.StringVar(&ver, "version", "", "")
	set.String("design", "", "")
	set.Parse(os.Args[1:])

	if err := codegen.CheckVersion(ver); err != nil {
		return nil, err
	}

	g := &Generator{OutDir: outDir, API: design.Design}

	return g.Generate()
}

// Generate produces the skeleton main.
func (g *Generator) Generate() (_ []string, err error) {
	go utils.Catch(nil, func() { g.Cleanup() })

	defer func() {
		if err != nil {
			g.Cleanup()
		}
	}()

	swagger, err := genswagger.New(g.API)
	if err != nil {
		return nil, err
	}

	swaggerDir := filepath.Join(g.OutDir, "swagger")
	os.RemoveAll(swaggerDir)
	if err = os.MkdirAll(swaggerDir, 0755); err != nil {
		return nil, err
	}
	g.genfiles = append(g.genfiles, swaggerDir)

	walk(swagger, appendKeys)

	for key = range keys {
		s, err := genswagger.New(g.API)
		if err != nil {
			return nil, err
		}
		extracted := walk(s, extract)
		suffix := "." + key

		// JSON
		rawJSON, err := json.Marshal(extracted)
		if err != nil {
			return nil, err
		}
		swaggerFile := filepath.Join(swaggerDir, fmt.Sprintf("swagger%s.json", suffix))
		if err := ioutil.WriteFile(swaggerFile, rawJSON, 0644); err != nil {
			return nil, err
		}
		g.genfiles = append(g.genfiles, swaggerFile)

		// YAML
		var yamlSource interface{}
		if err = json.Unmarshal(rawJSON, &yamlSource); err != nil {
			return nil, err
		}

		rawYAML, err := yaml.Marshal(yamlSource)
		if err != nil {
			return nil, err
		}
		swaggerFile = filepath.Join(swaggerDir, fmt.Sprintf("swagger%s.yaml", suffix))
		if err := ioutil.WriteFile(swaggerFile, rawYAML, 0644); err != nil {
			return nil, err
		}
		g.genfiles = append(g.genfiles, swaggerFile)
	}
	return g.genfiles, nil
}

// Cleanup removes all the files generated by this generator during the last invokation of Generate.
func (g *Generator) Cleanup() {
	for _, f := range g.genfiles {
		os.Remove(f)
	}
	g.genfiles = nil
}

var (
	keys = make(map[string]bool)
	key  string
)

func appendKeys(data *string) {
	var unmarshaled interface{}
	if err := json.Unmarshal([]byte(*data), &unmarshaled); err != nil {
		return
	}
	asserted := unmarshaled.(map[string]interface{})
	for k := range asserted {
		keys[k] = true
	}
}

func extract(data *string) {
	if *data == "" {
		return
	}
	var unmarshaled interface{}
	if err := json.Unmarshal([]byte(*data), &unmarshaled); err != nil {
		return
	}
	if len(keys) == 0 {
		return
	}
	asserted := unmarshaled.(map[string]interface{})
	value, ok := asserted[key]
	if ok != true {
		return
	}
	casted, ok := value.(string)
	if ok != true {
		return
	}
	*data = casted
}

func walk(s *genswagger.Swagger, function func(*string)) *genswagger.Swagger {
	if s.Info != nil {
		function(&s.Info.Description)
	}
	if s.Paths != nil {
		for _, v := range s.Paths {
			path, ok := v.(*genswagger.Path)
			if ok != true {
				continue
			}
			if path.Get != nil {
				function(&path.Get.Description)
			}
			if path.Put != nil {
				function(&path.Put.Description)
			}
			if path.Post != nil {
				function(&path.Post.Description)
			}
			if path.Delete != nil {
				function(&path.Delete.Description)
			}
			if path.Options != nil {
				function(&path.Options.Description)
			}
			if path.Head != nil {
				function(&path.Head.Description)
			}
			if path.Patch != nil {
				function(&path.Patch.Description)
			}
			if path.Parameters != nil {
				for _, parameter := range path.Parameters {
					function(&parameter.Description)
				}
			}
		}
	}
	if s.Parameters != nil {
		for _, parameter := range s.Parameters {
			function(&parameter.Description)
		}
	}
	if s.Responses != nil {
		for _, response := range s.Responses {
			function(&response.Description)
		}
	}
	if s.SecurityDefinitions != nil {
		for _, security := range s.SecurityDefinitions {
			function(&security.Description)
		}
	}
	if s.Tags != nil {
		for _, tag := range s.Tags {
			function(&tag.Description)
		}
	}
	if s.ExternalDocs != nil {
		function(&s.ExternalDocs.Description)
	}
	return s
}