Reflect optional output directory in the generated docs.go (#617)

gen/gen.go

@@ -13,9 +13,8 @@ import (
 	"text/template"
 	"time"
 
-	"github.com/go-openapi/spec"
-
 	"github.com/ghodss/yaml"
+	"github.com/go-openapi/spec"
 	"github.com/swaggo/swag"
 )
 
@@ -88,6 +87,7 @@ func (g *Gen) Build(config *Config) error {
 		return err
 	}
 
+	packageName := path.Base(config.OutputDir)
 	docFileName := path.Join(config.OutputDir, "docs.go")
 	jsonFileName := path.Join(config.OutputDir, "swagger.json")
 	yamlFileName := path.Join(config.OutputDir, "swagger.yaml")
@@ -105,7 +105,7 @@ func (g *Gen) Build(config *Config) error {
 
 	y, err := g.jsonToYAML(b)
 	if err != nil {
-		return fmt.Errorf("cannot covert json to yaml error: %s", err)
+		return fmt.Errorf("cannot convert json to yaml error: %s", err)
 	}
 
 	err = g.writeFile(y, yamlFileName)
@@ -114,14 +114,14 @@ func (g *Gen) Build(config *Config) error {
 	}
 
 	// Write doc
-	err = g.writeGoDoc(docs, swagger, config)
+	err = g.writeGoDoc(packageName, docs, swagger, config)
 	if err != nil {
 		return err
 	}
 
-	log.Printf("create docs.go at  %+v", docFileName)
-	log.Printf("create swagger.json at  %+v", jsonFileName)
-	log.Printf("create swagger.yaml at  %+v", yamlFileName)
+	log.Printf("create docs.go at %+v", docFileName)
+	log.Printf("create swagger.json at %+v", jsonFileName)
+	log.Printf("create swagger.yaml at %+v", yamlFileName)
 
 	return nil
 }
@@ -140,13 +140,12 @@ func (g *Gen) writeFile(b []byte, file string) error {
 func (g *Gen) formatSource(src []byte) []byte {
 	code, err := format.Source(src)
 	if err != nil {
-		code = src // Output the unformated code anyway
+		code = src // Output the unformatted code anyway
 	}
 	return code
 }
 
-func (g *Gen) writeGoDoc(output io.Writer, swagger *spec.Swagger, config *Config) error {
-
+func (g *Gen) writeGoDoc(packageName string, output io.Writer, swagger *spec.Swagger, config *Config) error {
 	generator, err := template.New("swagger_info").Funcs(template.FuncMap{
 		"printDoc": func(v string) string {
 			// Add schemes
@@ -202,6 +201,7 @@ func (g *Gen) writeGoDoc(output io.Writer, swagger *spec.Swagger, config *Config
 		GeneratedTime bool
 		Doc           string
 		Host          string
+		PackageName   string
 		BasePath      string
 		Schemes       []string
 		Title         string
@@ -212,6 +212,7 @@ func (g *Gen) writeGoDoc(output io.Writer, swagger *spec.Swagger, config *Config
 		GeneratedTime: config.GeneratedTime,
 		Doc:           string(buf),
 		Host:          swagger.Host,
+		PackageName:   packageName,
 		BasePath:      swagger.BasePath,
 		Schemes:       swagger.Schemes,
 		Title:         swagger.Info.Title,
@@ -233,7 +234,7 @@ var packageTemplate = `// GENERATED BY THE COMMAND ABOVE; DO NOT EDIT
 // This file was generated by swaggo/swag{{ if .GeneratedTime }} at
 // {{ .Timestamp }}{{ end }}
 
-package docs
+package {{.PackageName}}
 
 import (
 	"bytes"

gen/gen_test.go

@@ -2,13 +2,13 @@ package gen
 
 import (
 	"errors"
-	"github.com/go-openapi/spec"
 	"os"
 	"os/exec"
 	"path"
 	"path/filepath"
 	"testing"
 
+	"github.com/go-openapi/spec"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -262,11 +262,11 @@ fmt.Print("Helo world")
 	assert.NotEqual(t, []byte(src2), res, "Should return fmt code")
 }
 
-type mocWriter struct {
+type mockWriter struct {
 	hook func([]byte)
 }
 
-func (w *mocWriter) Write(data []byte) (int, error) {
+func (w *mockWriter) Write(data []byte) (int, error) {
 	if w.hook != nil {
 		w.hook(data)
 	}
@@ -279,7 +279,7 @@ func TestGen_writeGoDoc(t *testing.T) {
 	swapTemplate := packageTemplate
 
 	packageTemplate = `{{{`
-	err := gen.writeGoDoc(nil, nil, &Config{})
+	err := gen.writeGoDoc("docs", nil, nil, &Config{})
 	assert.Error(t, err)
 
 	packageTemplate = `{{.Data}}`
@@ -289,21 +289,23 @@ func TestGen_writeGoDoc(t *testing.T) {
 			Info: &spec.Info{},
 		},
 	}
-	err = gen.writeGoDoc(&mocWriter{}, swagger, &Config{})
+	err = gen.writeGoDoc("docs", &mockWriter{}, swagger, &Config{})
 	assert.Error(t, err)
 
 	packageTemplate = `{{ if .GeneratedTime }}Fake Time{{ end }}`
-	err = gen.writeGoDoc(&mocWriter{
-		hook: func(data []byte) {
-			assert.Equal(t, "Fake Time", string(data))
-		},
-	}, swagger, &Config{GeneratedTime: true})
+	err = gen.writeGoDoc("docs",
+		&mockWriter{
+			hook: func(data []byte) {
+				assert.Equal(t, "Fake Time", string(data))
+			},
+		}, swagger, &Config{GeneratedTime: true})
 	assert.NoError(t, err)
-	err = gen.writeGoDoc(&mocWriter{
-		hook: func(data []byte) {
-			assert.Equal(t, "", string(data))
-		},
-	}, swagger, &Config{GeneratedTime: false})
+	err = gen.writeGoDoc("docs",
+		&mockWriter{
+			hook: func(data []byte) {
+				assert.Equal(t, "", string(data))
+			},
+		}, swagger, &Config{GeneratedTime: false})
 	assert.NoError(t, err)
 
 	packageTemplate = swapTemplate