xrand.go

Documentation: oss.terrastruct.com/util-go/xrand

     1  // Package xrand provides helpers for generating useful random values.
     2  package xrand
     3  
     4  import (
     5  	"crypto/rand"
     6  	"encoding/base64"
     7  	"fmt"
     8  	"io"
     9  	math_rand "math/rand"
    10  	"strings"
    11  	"time"
    12  	"unicode/utf8"
    13  )
    14  
    15  func init() {
    16  	math_rand.Seed(time.Now().UnixNano())
    17  }
    18  
    19  // 0 is a valid unicode codepoint but certain programs
    20  // won't allow 0 to avoid clashing with c strings. e.g. pq.
    21  func randRune() rune {
    22  	for {
    23  		if Bool() {
    24  			// Generate plain ASCII half the time.
    25  			if math_rand.Int31n(100) == 0 {
    26  				// Generate newline 1% of the time.
    27  				return '\n'
    28  			}
    29  			return math_rand.Int31n(128) + 1
    30  		}
    31  		r := math_rand.Int31n(utf8.MaxRune+1) + 1
    32  		if utf8.ValidRune(r) {
    33  			return r
    34  		}
    35  	}
    36  }
    37  
    38  func String(n int, exclude []rune) string {
    39  	var b strings.Builder
    40  	for i := 0; i < n; i++ {
    41  		r := randRune()
    42  		excluded := false
    43  		for _, xr := range exclude {
    44  			if r == xr {
    45  				excluded = true
    46  				break
    47  			}
    48  		}
    49  		if excluded {
    50  			i--
    51  			continue
    52  		}
    53  		b.WriteRune(r)
    54  	}
    55  	return b.String()
    56  }
    57  
    58  func Bool() bool {
    59  	return math_rand.Intn(2) == 0
    60  }
    61  
    62  func Base64(n int) string {
    63  	var b strings.Builder
    64  	b.Grow(n)
    65  	wc := base64.NewEncoder(base64.URLEncoding, &b)
    66  
    67  	_, err := io.CopyN(wc, rand.Reader, int64(n))
    68  	if err != nil {
    69  		panic(fmt.Sprintf("error encoding base64: %v", err))
    70  	}
    71  
    72  	err = wc.Close()
    73  	if err != nil {
    74  		panic(fmt.Sprintf("error encoding base64: %v", err))
    75  	}
    76  
    77  	return b.String()
    78  }
    79
View as plain text