github.com/cellofellow/gopkg@v0.0.0-20140722061823-eec0544a62ad/database/redis/format.go (about)

     1  // Copyright 2014 <chaishushan{AT}gmail.com>. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package redis
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"reflect"
    11  	"strconv"
    12  )
    13  
    14  var delim []byte = []byte{'\r', '\n'}
    15  
    16  type request struct {
    17  	cmd  string
    18  	args []interface{}
    19  }
    20  
    21  // formatArg formats the given argument to a Redis-styled argument byte slice.
    22  func formatArg(v interface{}) []byte {
    23  	var b, bs []byte
    24  
    25  	switch vt := v.(type) {
    26  	case []byte:
    27  		bs = vt
    28  	case string:
    29  		bs = []byte(vt)
    30  	case bool:
    31  		if vt {
    32  			bs = []byte{'1'}
    33  		} else {
    34  			bs = []byte{'0'}
    35  		}
    36  	case nil:
    37  		// empty byte slice
    38  	case int:
    39  		bs = []byte(strconv.Itoa(vt))
    40  	case int8:
    41  		bs = []byte(strconv.FormatInt(int64(vt), 10))
    42  	case int16:
    43  		bs = []byte(strconv.FormatInt(int64(vt), 10))
    44  	case int32:
    45  		bs = []byte(strconv.FormatInt(int64(vt), 10))
    46  	case int64:
    47  		bs = []byte(strconv.FormatInt(vt, 10))
    48  	case uint:
    49  		bs = []byte(strconv.FormatUint(uint64(vt), 10))
    50  	case uint8:
    51  		bs = []byte(strconv.FormatUint(uint64(vt), 10))
    52  	case uint16:
    53  		bs = []byte(strconv.FormatUint(uint64(vt), 10))
    54  	case uint32:
    55  		bs = []byte(strconv.FormatUint(uint64(vt), 10))
    56  	case uint64:
    57  		bs = []byte(strconv.FormatUint(vt, 10))
    58  	default:
    59  		// Fallback to reflect-based.
    60  		switch reflect.TypeOf(vt).Kind() {
    61  		case reflect.Slice:
    62  			rv := reflect.ValueOf(vt)
    63  			for i := 0; i < rv.Len(); i++ {
    64  				bs = append(bs, formatArg(rv.Index(i).Interface())...)
    65  			}
    66  
    67  			return bs
    68  		case reflect.Map:
    69  			rv := reflect.ValueOf(vt)
    70  			keys := rv.MapKeys()
    71  			for _, k := range keys {
    72  				bs = append(bs, formatArg(k.Interface())...)
    73  				bs = append(bs, formatArg(rv.MapIndex(k).Interface())...)
    74  			}
    75  
    76  			return bs
    77  		default:
    78  			var buf bytes.Buffer
    79  
    80  			fmt.Fprint(&buf, v)
    81  			bs = buf.Bytes()
    82  		}
    83  	}
    84  
    85  	b = append(b, '$')
    86  	b = append(b, []byte(strconv.Itoa(len(bs)))...)
    87  	b = append(b, delim...)
    88  	b = append(b, bs...)
    89  	b = append(b, delim...)
    90  	return b
    91  }
    92  
    93  // createRequest creates a request string from the given requests.
    94  func createRequest(requests ...*request) []byte {
    95  	var total []byte
    96  
    97  	for _, req := range requests {
    98  		var s []byte
    99  
   100  		// Calculate number of arguments.
   101  		argsLen := 1
   102  		for _, arg := range req.args {
   103  			kind := reflect.TypeOf(arg).Kind()
   104  			switch kind {
   105  			case reflect.Slice:
   106  				argsLen += reflect.ValueOf(arg).Len()
   107  			case reflect.Map:
   108  				argsLen += reflect.ValueOf(arg).Len() * 2
   109  			default:
   110  				argsLen++
   111  			}
   112  		}
   113  
   114  		// number of arguments
   115  		s = append(s, '*')
   116  		s = append(s, []byte(strconv.Itoa(argsLen))...)
   117  		s = append(s, delim...)
   118  
   119  		// command
   120  		s = append(s, '$')
   121  		s = append(s, []byte(strconv.Itoa(len(req.cmd)))...)
   122  		s = append(s, delim...)
   123  		s = append(s, []byte(req.cmd)...)
   124  		s = append(s, delim...)
   125  
   126  		// arguments
   127  		for _, arg := range req.args {
   128  			s = append(s, formatArg(arg)...)
   129  		}
   130  
   131  		total = append(total, s...)
   132  	}
   133  
   134  	return total
   135  }