Add ability to generate random string from any charset

README.md

@@ -39,7 +39,3 @@ From the `math/rand` [docs](https://golang.org/pkg/math/rand/):
 ## Credits
 
 I claim no credit for the code here. It's originally from user icza in https://stackoverflow.com/a/31832326/1161743.
-
-## Possible future features
-
-1. Allow custom bytes, e.g. to include special characters

grand.go

@@ -4,17 +4,34 @@ import "math/rand"
 
 // Credits to icza of https://stackoverflow.com/a/31832326/1161743 for the original code.
 
-const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+// Character sets for random string generation
+const (
+	CharSetEnglishAlphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+	CharSetBase62          = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+)
+
+var defaultLetterBytes = CharSetEnglishAlphabet
 
 const (
 	letterIdxBits = 6                    // 6 bits to represent a letter index
 	letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
 	letterIdxMax  = 63 / letterIdxBits   // # of letter indices fitting in 63 bits
 )
 
-// GenerateRandomString generates a length-n random string from the character set "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ".
-// Be sure to seed the random function using rand.Seed to initialize the generator first.
-func GenerateRandomString(n int) string {
+// Generator is the instance that generates random strings from the given charset
+type Generator struct {
+	charset string
+}
+
+// NewGenerator returns a new Generator instance
+func NewGenerator(charset string) *Generator {
+	return &Generator{
+		charset: charset,
+	}
+}
+
+// GenerateRandomString generates a length-n random string from the given character set
+func (g *Generator) GenerateRandomString(n int) string {
 	b := make([]byte, n)
 
 	// A rand.Int63() generates 63 random bits, enough for letterIdxMax letters!
@@ -23,8 +40,8 @@ func GenerateRandomString(n int) string {
 			cache, remain = rand.Int63(), letterIdxMax
 		}
 
-		if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
-			b[i] = letterBytes[idx]
+		if idx := int(cache & letterIdxMask); idx < len(g.charset) {
+			b[i] = g.charset[idx]
 			i--
 		}
 		cache >>= letterIdxBits
@@ -33,3 +50,9 @@ func GenerateRandomString(n int) string {
 
 	return string(b)
 }
+
+// GenerateRandomString generates a length-n random string from the default character set (English alphabets).
+// Be sure to seed the random function using rand.Seed to initialize the generator first.
+func GenerateRandomString(n int) string {
+	return NewGenerator(CharSetEnglishAlphabet).GenerateRandomString(n)
+}

grand_test.go

@@ -2,6 +2,7 @@ package grand
 
 import (
 	"math/rand"
+	"regexp"
 	"testing"
 	"time"
 )
@@ -27,3 +28,25 @@ func TestGeneratesDifferentStrings(t *testing.T) {
 		t.Fatal()
 	}
 }
+
+func TestSetDefaultCharSet(t *testing.T) {
+	randStr1 := GenerateRandomString(20)
+
+	// There should not be any numbers in the string
+	res := regexp.MustCompile("\\d").FindAllString(randStr1, -1)
+	if len(res) != 0 {
+		t.Fatal()
+	}
+}
+
+func TestSetDifferentCharSet(t *testing.T) {
+	gen := NewGenerator(CharSetBase62)
+
+	randStr1 := gen.GenerateRandomString(20)
+
+	// There should be at least one number in generated string
+	res := regexp.MustCompile("\\d").FindAllString(randStr1, -1)
+	if len(res) == 0 {
+		t.Fatal()
+	}
+}