Support multiple assets

fs/fs.go

@@ -30,7 +30,7 @@ import (
 	"time"
 )
 
-var zipData string
+var zipData = map[string]string{}
 
 // file holds unzipped read-only file contents and file metadata.
 type file struct {
@@ -47,16 +47,29 @@ type statikFS struct {
 // Register registers zip contents data, later used to initialize
 // the statik file system.
 func Register(data string) {
-	zipData = data
+	RegisterWithName("default", data)
 }
 
-// New creates a new file system with the registered zip contents data.
+// RegisterWithName registers zip contents data and set asset namespace,
+// later used to initialize the statik file system.
+func RegisterWithName(assetNamespace string, data string) {
+	zipData[assetNamespace] = data
+}
+
+// New creates a new file system with the default registered zip contents data.
 // It unzips all files and stores them in an in-memory map.
 func New() (http.FileSystem, error) {
-	if zipData == "" {
+	return NewWithName("default")
+}
+
+// NewWithName creates a new file system with the registered zip contents data.
+// It unzips all files and stores them in an in-memory map.
+func NewWithName(assetNamespace string) (http.FileSystem, error) {
+	asset, ok := zipData[assetNamespace]
+	if !ok {
 		return nil, errors.New("statik/fs: no zip data registered")
 	}
-	zipReader, err := zip.NewReader(strings.NewReader(zipData), int64(len(zipData)))
+	zipReader, err := zip.NewReader(strings.NewReader(asset), int64(len(asset)))
 	if err != nil {
 		return nil, err
 	}

fs/fs_test.go

@@ -16,6 +16,7 @@ package fs
 import (
 	"archive/zip"
 	"bytes"
+	"errors"
 	"io"
 	"io/ioutil"
 	"os"
@@ -39,6 +40,70 @@ type wantFile struct {
 	err     error
 }
 
+func TestRegisterWithName(t *testing.T) {
+	tests := []struct {
+		description string
+		assetName   string
+		zipData     string
+		condition   func() error
+	}{
+		{
+			description: "RegisterWithName() should set zipData with assetName to be key",
+			assetName:   "file",
+			zipData:     "file test",
+			condition: func() error {
+				data, ok := zipData["file"]
+				if !ok {
+					return errors.New("fail to register zipData")
+				}
+				if data != "file test" {
+					return errors.New("fail to register zipData[\"file\"]")
+				}
+				return nil
+			},
+		},
+		{
+			description: "zipData[\"default\"] should be able to open by Open()",
+			assetName:   "default",
+			zipData:     mustZipTree("../testdata/file"),
+			condition: func() error {
+				fs, err := New()
+				if err != nil {
+					return err
+				}
+				if _, err := fs.Open("/file.txt"); err != nil {
+					return err
+				}
+				return nil
+			},
+		},
+		{
+			description: "zipData[\"foo\"] should be able to open by OpenWithName(\"foo\")",
+			assetName:   "foo",
+			zipData:     mustZipTree("../testdata/file"),
+			condition: func() error {
+				fs, err := NewWithName("foo")
+				if err != nil {
+					return err
+				}
+				if _, err := fs.Open("/file.txt"); err != nil {
+					return err
+				}
+				return nil
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.description, func(t *testing.T) {
+			RegisterWithName(tc.assetName, tc.zipData)
+			if err := tc.condition(); err != nil {
+				t.Error(err)
+			}
+		})
+		delete(zipData, tc.assetName)
+	}
+}
+
 func TestOpen(t *testing.T) {
 	fileTxtHeader := mustFileHeader("../testdata/file/file.txt")
 	pixelGifHeader := mustFileHeader("../testdata/image/pixel.gif")

statik.go

@@ -28,6 +28,7 @@ import (
 	"path/filepath"
 	"strings"
 	"time"
+	"unicode"
 )
 
 const (
@@ -44,6 +45,7 @@ var (
 	flagForce      = flag.Bool("f", false, "Overwrite destination file if it already exists.")
 	flagTags       = flag.String("tags", "", "Write build constraint tags")
 	flagPkg        = flag.String("p", "statik", "Name of the generated package")
+	flagAstName    = flag.String("a", "default", "Set static assets name")
 	flagPkgCmt     = flag.String("c", "Package statik contains static assets.", "The package comment. An empty value disables this comment.\n")
 )
 
@@ -190,6 +192,12 @@ func generateSource(srcPath string) (file *os.File, err error) {
 		comment = "\n" + commentLines(*flagPkgCmt)
 	}
 
+	// e.g.)
+	// assetNamespaceIdentify is "AbcDeF_G"
+	// when assetNamespace is "abc de f-g"
+	assetNamespace := *flagAstName
+	assetNamespaceIdentify := toSymbolSafe(assetNamespace)
+
 	// then embed it as a quoted string
 	var qb bytes.Buffer
 	fmt.Fprintf(&qb, `// Code generated by statik. DO NOT EDIT.
@@ -200,13 +208,20 @@ import (
 	"github.com/rakyll/statik/fs"
 )
 
+`, tags, comment, namePackage)
+	if assetNamespace != "default" {
+		fmt.Fprintf(&qb, `
+const %s = "%s" // static asset namespace
+`, assetNamespaceIdentify, assetNamespace)
+	}
+	fmt.Fprint(&qb, `
 func init() {
-	data := "`, tags, comment, namePackage)
+	data := "`)
 	FprintZipData(&qb, buffer.Bytes())
-	fmt.Fprint(&qb, `"
-	fs.Register(data)
+	fmt.Fprintf(&qb, `"
+	fs.RegisterWithName("%s", data)
 }
-`)
+`, assetNamespace)
 
 	if err = ioutil.WriteFile(f.Name(), qb.Bytes(), 0644); err != nil {
 		return
@@ -248,3 +263,31 @@ func exitWithError(err error) {
 	fmt.Println(err)
 	os.Exit(1)
 }
+
+// convert src to symbol safe string with upper camel case
+func toSymbolSafe(str string) string {
+	isBeforeRuneNoGeneralCase := false
+	replace := func(r rune) rune {
+		if unicode.IsLetter(r) {
+			if isBeforeRuneNoGeneralCase {
+				isBeforeRuneNoGeneralCase = true
+				return r
+			} else {
+				isBeforeRuneNoGeneralCase = true
+				return unicode.ToTitle(r)
+			}
+		} else if unicode.IsDigit(r) {
+			if isBeforeRuneNoGeneralCase {
+				isBeforeRuneNoGeneralCase = true
+				return r
+			} else {
+				isBeforeRuneNoGeneralCase = false
+				return -1
+			}
+		} else {
+			isBeforeRuneNoGeneralCase = false
+			return -1
+		}
+	}
+	return strings.TrimSpace(strings.Map(replace, str))
+}

statik_test.go

@@ -0,0 +1,21 @@
+package main
+
+import "testing"
+
+func TestToSymbolSafe(t *testing.T) {
+	testCase := [][]string{
+		{"abc", "Abc"},
+		{"_abc", "Abc"},
+		{"3abc", "Abc"},
+		{"abc3", "Abc3"},
+		{"/abc", "Abc"},
+		{"abc abc", "AbcAbc"},
+	}
+	for i, test := range testCase {
+		got := toSymbolSafe(test[0])
+		wont := test[1]
+		if got != wont {
+			t.Errorf("#%02d toSymbolSafe(%s) => %s != %s", i, test[0], got, wont)
+		}
+	}
+}