transform.go

Documentation: github.com/thoas/go-funk

     1  package funk
     2  
     3  import (
     4  	"fmt"
     5  	"math/rand"
     6  	"reflect"
     7  	"strings"
     8  )
     9  
    10  // Chunk creates an array of elements split into groups with the length of size.
    11  // If array can't be split evenly, the final chunk will be
    12  // the remaining element.
    13  func Chunk(arr interface{}, size int) interface{} {
    14  	if !IsIteratee(arr) {
    15  		panic("First parameter must be neither array nor slice")
    16  	}
    17  
    18  	if size == 0 {
    19  		return arr
    20  	}
    21  
    22  	arrValue := reflect.ValueOf(arr)
    23  
    24  	arrType := arrValue.Type()
    25  
    26  	resultSliceType := reflect.SliceOf(arrType)
    27  
    28  	// Initialize final result slice which will contains slice
    29  	resultSlice := reflect.MakeSlice(resultSliceType, 0, 0)
    30  
    31  	itemType := arrType.Elem()
    32  
    33  	var itemSlice reflect.Value
    34  
    35  	itemSliceType := reflect.SliceOf(itemType)
    36  
    37  	length := arrValue.Len()
    38  
    39  	for i := 0; i < length; i++ {
    40  		if i%size == 0 || i == 0 {
    41  			if itemSlice.Kind() != reflect.Invalid {
    42  				resultSlice = reflect.Append(resultSlice, itemSlice)
    43  			}
    44  
    45  			itemSlice = reflect.MakeSlice(itemSliceType, 0, 0)
    46  		}
    47  
    48  		itemSlice = reflect.Append(itemSlice, arrValue.Index(i))
    49  
    50  		if i == length-1 {
    51  			resultSlice = reflect.Append(resultSlice, itemSlice)
    52  		}
    53  	}
    54  
    55  	return resultSlice.Interface()
    56  }
    57  
    58  // ToMap transforms a slice of instances to a Map.
    59  // []*Foo => Map<int, *Foo>
    60  func ToMap(in interface{}, pivot string) interface{} {
    61  	value := reflect.ValueOf(in)
    62  
    63  	// input value must be a slice
    64  	if value.Kind() != reflect.Slice {
    65  		panic(fmt.Sprintf("%v must be a slice", in))
    66  	}
    67  
    68  	inType := value.Type()
    69  
    70  	structType := inType.Elem()
    71  
    72  	// retrieve the struct in the slice to deduce key type
    73  	if structType.Kind() == reflect.Ptr {
    74  		structType = structType.Elem()
    75  	}
    76  
    77  	field, _ := structType.FieldByName(pivot)
    78  
    79  	// value of the map will be the input type
    80  	collectionType := reflect.MapOf(field.Type, inType.Elem())
    81  
    82  	// create a map from scratch
    83  	collection := reflect.MakeMap(collectionType)
    84  
    85  	for i := 0; i < value.Len(); i++ {
    86  		instance := value.Index(i)
    87  		var field reflect.Value
    88  
    89  		if instance.Kind() == reflect.Ptr {
    90  			field = instance.Elem().FieldByName(pivot)
    91  		} else {
    92  			field = instance.FieldByName(pivot)
    93  		}
    94  
    95  		collection.SetMapIndex(field, instance)
    96  	}
    97  
    98  	return collection.Interface()
    99  }
   100  
   101  func mapSlice(arrValue reflect.Value, funcValue reflect.Value) reflect.Value {
   102  	funcType := funcValue.Type()
   103  
   104  	if funcType.NumIn() != 1 || funcType.NumOut() == 0 || funcType.NumOut() > 2 {
   105  		panic("Map function with an array must have one parameter and must return one or two parameters")
   106  	}
   107  
   108  	arrElemType := arrValue.Type().Elem()
   109  
   110  	// Checking whether element type is convertible to function's first argument's type.
   111  	if !arrElemType.ConvertibleTo(funcType.In(0)) {
   112  		panic("Map function's argument is not compatible with type of array.")
   113  	}
   114  
   115  	if funcType.NumOut() == 1 {
   116  		// Get slice type corresponding to function's return value's type.
   117  		resultSliceType := reflect.SliceOf(funcType.Out(0))
   118  
   119  		// MakeSlice takes a slice kind type, and makes a slice.
   120  		resultSlice := reflect.MakeSlice(resultSliceType, 0, 0)
   121  
   122  		for i := 0; i < arrValue.Len(); i++ {
   123  			result := funcValue.Call([]reflect.Value{arrValue.Index(i)})[0]
   124  
   125  			resultSlice = reflect.Append(resultSlice, result)
   126  		}
   127  
   128  		return resultSlice
   129  	}
   130  
   131  	if funcType.NumOut() == 2 {
   132  		// value of the map will be the input type
   133  		collectionType := reflect.MapOf(funcType.Out(0), funcType.Out(1))
   134  
   135  		// create a map from scratch
   136  		collection := reflect.MakeMap(collectionType)
   137  
   138  		for i := 0; i < arrValue.Len(); i++ {
   139  			results := funcValue.Call([]reflect.Value{arrValue.Index(i)})
   140  
   141  			collection.SetMapIndex(results[0], results[1])
   142  		}
   143  
   144  		return collection
   145  	}
   146  
   147  	return reflect.Value{}
   148  }
   149  
   150  func mapMap(arrValue reflect.Value, funcValue reflect.Value) reflect.Value {
   151  	funcType := funcValue.Type()
   152  
   153  	if funcType.NumIn() != 2 || funcType.NumOut() == 0 || funcType.NumOut() > 2 {
   154  		panic("Map function with a map must have two parameters and must return one or two parameters")
   155  	}
   156  
   157  	// Only one returned parameter, should be a slice
   158  	if funcType.NumOut() == 1 {
   159  		// Get slice type corresponding to function's return value's type.
   160  		resultSliceType := reflect.SliceOf(funcType.Out(0))
   161  
   162  		// MakeSlice takes a slice kind type, and makes a slice.
   163  		resultSlice := reflect.MakeSlice(resultSliceType, 0, 0)
   164  
   165  		for _, key := range arrValue.MapKeys() {
   166  			results := funcValue.Call([]reflect.Value{key, arrValue.MapIndex(key)})
   167  
   168  			result := results[0]
   169  
   170  			resultSlice = reflect.Append(resultSlice, result)
   171  		}
   172  
   173  		return resultSlice
   174  	}
   175  
   176  	// two parameters, should be a map
   177  	if funcType.NumOut() == 2 {
   178  		// value of the map will be the input type
   179  		collectionType := reflect.MapOf(funcType.Out(0), funcType.Out(1))
   180  
   181  		// create a map from scratch
   182  		collection := reflect.MakeMap(collectionType)
   183  
   184  		for _, key := range arrValue.MapKeys() {
   185  			results := funcValue.Call([]reflect.Value{key, arrValue.MapIndex(key)})
   186  
   187  			collection.SetMapIndex(results[0], results[1])
   188  
   189  		}
   190  
   191  		return collection
   192  	}
   193  
   194  	return reflect.Value{}
   195  }
   196  
   197  // Map manipulates an iteratee and transforms it to another type.
   198  func Map(arr interface{}, mapFunc interface{}) interface{} {
   199  	result := mapFn(arr, mapFunc, "Map")
   200  
   201  	if result.IsValid() {
   202  		return result.Interface()
   203  	}
   204  
   205  	return nil
   206  }
   207  
   208  func mapFn(arr interface{}, mapFunc interface{}, funcName string) reflect.Value {
   209  	if !IsIteratee(arr) {
   210  		panic("First parameter must be an iteratee")
   211  	}
   212  
   213  	if !IsFunction(mapFunc) {
   214  		panic("Second argument must be function")
   215  	}
   216  
   217  	var (
   218  		funcValue = reflect.ValueOf(mapFunc)
   219  		arrValue  = reflect.ValueOf(arr)
   220  		arrType   = arrValue.Type()
   221  	)
   222  
   223  	kind := arrType.Kind()
   224  
   225  	if kind == reflect.Slice || kind == reflect.Array {
   226  		return mapSlice(arrValue, funcValue)
   227  	} else if kind == reflect.Map {
   228  		return mapMap(arrValue, funcValue)
   229  	}
   230  
   231  	panic(fmt.Sprintf("Type %s is not supported by "+funcName, arrType.String()))
   232  }
   233  
   234  // FlatMap manipulates an iteratee and transforms it to a flattened collection of another type.
   235  func FlatMap(arr interface{}, mapFunc interface{}) interface{} {
   236  	result := mapFn(arr, mapFunc, "FlatMap")
   237  
   238  	if result.IsValid() {
   239  		return flatten(result).Interface()
   240  	}
   241  
   242  	return nil
   243  }
   244  
   245  // Flatten flattens a two-dimensional array.
   246  func Flatten(out interface{}) interface{} {
   247  	return flatten(reflect.ValueOf(out)).Interface()
   248  }
   249  
   250  func flatten(value reflect.Value) reflect.Value {
   251  	sliceType := value.Type()
   252  
   253  	if (value.Kind() != reflect.Slice && value.Kind() != reflect.Array) ||
   254  		(sliceType.Elem().Kind() != reflect.Slice && sliceType.Elem().Kind() != reflect.Array) {
   255  		panic("Argument must be an array or slice of at least two dimensions")
   256  	}
   257  
   258  	resultSliceType := sliceType.Elem().Elem()
   259  
   260  	resultSlice := reflect.MakeSlice(reflect.SliceOf(resultSliceType), 0, 0)
   261  
   262  	length := value.Len()
   263  
   264  	for i := 0; i < length; i++ {
   265  		item := value.Index(i)
   266  
   267  		resultSlice = reflect.AppendSlice(resultSlice, item)
   268  	}
   269  
   270  	return resultSlice
   271  }
   272  
   273  // FlattenDeep recursively flattens array.
   274  func FlattenDeep(out interface{}) interface{} {
   275  	return flattenDeep(reflect.ValueOf(out)).Interface()
   276  }
   277  
   278  func flattenDeep(value reflect.Value) reflect.Value {
   279  	sliceType := sliceElem(value.Type())
   280  
   281  	resultSlice := reflect.MakeSlice(reflect.SliceOf(sliceType), 0, 0)
   282  
   283  	return flattenRecursive(value, resultSlice)
   284  }
   285  
   286  func flattenRecursive(value reflect.Value, result reflect.Value) reflect.Value {
   287  	length := value.Len()
   288  
   289  	for i := 0; i < length; i++ {
   290  		item := value.Index(i)
   291  		kind := item.Kind()
   292  
   293  		if kind == reflect.Slice || kind == reflect.Array {
   294  			result = flattenRecursive(item, result)
   295  		} else {
   296  			result = reflect.Append(result, item)
   297  		}
   298  	}
   299  
   300  	return result
   301  }
   302  
   303  // Shuffle creates an array of shuffled values
   304  func Shuffle(in interface{}) interface{} {
   305  	value := reflect.ValueOf(in)
   306  	valueType := value.Type()
   307  
   308  	kind := value.Kind()
   309  
   310  	if kind == reflect.Array || kind == reflect.Slice {
   311  		length := value.Len()
   312  
   313  		resultSlice := makeSlice(value, length)
   314  
   315  		for i, v := range rand.Perm(length) {
   316  			resultSlice.Index(i).Set(value.Index(v))
   317  		}
   318  
   319  		return resultSlice.Interface()
   320  	}
   321  
   322  	panic(fmt.Sprintf("Type %s is not supported by Shuffle", valueType.String()))
   323  }
   324  
   325  // Reverse transforms an array the first element will become the last,
   326  // the second element will become the second to last, etc.
   327  func Reverse(in interface{}) interface{} {
   328  	value := reflect.ValueOf(in)
   329  	valueType := value.Type()
   330  
   331  	kind := value.Kind()
   332  
   333  	if kind == reflect.String {
   334  		return ReverseString(in.(string))
   335  	}
   336  
   337  	if kind == reflect.Array || kind == reflect.Slice {
   338  		length := value.Len()
   339  
   340  		resultSlice := makeSlice(value, length)
   341  
   342  		j := 0
   343  		for i := length - 1; i >= 0; i-- {
   344  			resultSlice.Index(j).Set(value.Index(i))
   345  			j++
   346  		}
   347  
   348  		return resultSlice.Interface()
   349  	}
   350  
   351  	panic(fmt.Sprintf("Type %s is not supported by Reverse", valueType.String()))
   352  }
   353  
   354  // Uniq creates an array with unique values.
   355  func Uniq(in interface{}) interface{} {
   356  	value := reflect.ValueOf(in)
   357  	valueType := value.Type()
   358  
   359  	kind := value.Kind()
   360  
   361  	if kind == reflect.Array || kind == reflect.Slice {
   362  		length := value.Len()
   363  
   364  		result := makeSlice(value, 0)
   365  
   366  		seen := make(map[interface{}]bool, length)
   367  		j := 0
   368  
   369  		for i := 0; i < length; i++ {
   370  			val := value.Index(i)
   371  			v := val.Interface()
   372  
   373  			if _, ok := seen[v]; ok {
   374  				continue
   375  			}
   376  
   377  			seen[v] = true
   378  			result = reflect.Append(result, val)
   379  			j++
   380  		}
   381  
   382  		return result.Interface()
   383  	}
   384  
   385  	panic(fmt.Sprintf("Type %s is not supported by Uniq", valueType.String()))
   386  }
   387  
   388  // ConvertSlice converts a slice type to another,
   389  // a perfect example would be to convert a slice of struct to a slice of interface.
   390  func ConvertSlice(in interface{}, out interface{}) {
   391  	srcValue := reflect.ValueOf(in)
   392  
   393  	dstValue := reflect.ValueOf(out)
   394  
   395  	if dstValue.Kind() != reflect.Ptr {
   396  		panic("Second argument must be a pointer")
   397  	}
   398  
   399  	dstValue = dstValue.Elem()
   400  
   401  	if srcValue.Kind() != reflect.Slice && srcValue.Kind() != reflect.Array {
   402  		panic("First argument must be an array or slice")
   403  	}
   404  
   405  	if dstValue.Kind() != reflect.Slice && dstValue.Kind() != reflect.Array {
   406  		panic("Second argument must be an array or slice")
   407  	}
   408  
   409  	// returns value that points to dstValue
   410  	direct := reflect.Indirect(dstValue)
   411  
   412  	length := srcValue.Len()
   413  
   414  	for i := 0; i < length; i++ {
   415  		dstValue = reflect.Append(dstValue, srcValue.Index(i))
   416  	}
   417  
   418  	direct.Set(dstValue)
   419  }
   420  
   421  // Drop creates an array/slice with `n` elements dropped from the beginning.
   422  func Drop(in interface{}, n int) interface{} {
   423  	value := reflect.ValueOf(in)
   424  	valueType := value.Type()
   425  
   426  	kind := value.Kind()
   427  
   428  	if kind == reflect.Array || kind == reflect.Slice {
   429  		length := value.Len()
   430  
   431  		resultSlice := makeSlice(value, length-n)
   432  
   433  		j := 0
   434  		for i := n; i < length; i++ {
   435  			resultSlice.Index(j).Set(value.Index(i))
   436  			j++
   437  		}
   438  
   439  		return resultSlice.Interface()
   440  
   441  	}
   442  
   443  	panic(fmt.Sprintf("Type %s is not supported by Drop", valueType.String()))
   444  }
   445  
   446  // Prune returns a copy of "in" that only contains fields in "paths"
   447  // which are looked up using struct field name.
   448  // For lookup paths by field tag instead, use funk.PruneByTag()
   449  func Prune(in interface{}, paths []string) (interface{}, error) {
   450  	return pruneByTag(in, paths, nil /*tag*/)
   451  }
   452  
   453  // pruneByTag returns a copy of "in" that only contains fields in "paths"
   454  // which are looked up using struct field Tag "tag".
   455  func PruneByTag(in interface{}, paths []string, tag string) (interface{}, error) {
   456  	return pruneByTag(in, paths, &tag)
   457  }
   458  
   459  // pruneByTag returns a copy of "in" that only contains fields in "paths"
   460  // which are looked up using struct field Tag "tag". If tag is nil,
   461  // traverse paths using struct field name
   462  func pruneByTag(in interface{}, paths []string, tag *string) (interface{}, error) {
   463  
   464  	inValue := reflect.ValueOf(in)
   465  
   466  	ret := reflect.New(inValue.Type()).Elem()
   467  
   468  	for _, path := range paths {
   469  		parts := strings.Split(path, ".")
   470  		if err := prune(inValue, ret, parts, tag); err != nil {
   471  			return nil, err
   472  		}
   473  	}
   474  	return ret.Interface(), nil
   475  }
   476  
   477  func prune(inValue reflect.Value, ret reflect.Value, parts []string, tag *string) error {
   478  
   479  	if len(parts) == 0 {
   480  		// we reached the location that ret needs to hold inValue
   481  		// Note: The value at the end of the path is not copied, maybe we need to change.
   482  		// ret and the original data holds the same reference to this value
   483  		ret.Set(inValue)
   484  		return nil
   485  	}
   486  
   487  	inKind := inValue.Kind()
   488  
   489  	switch inKind {
   490  	case reflect.Ptr:
   491  		if inValue.IsNil() {
   492  			// TODO validate
   493  			return nil
   494  		}
   495  		if ret.IsNil() {
   496  			// init ret and go to next level
   497  			ret.Set(reflect.New(inValue.Type().Elem()))
   498  		}
   499  		return prune(inValue.Elem(), ret.Elem(), parts, tag)
   500  	case reflect.Struct:
   501  		part := parts[0]
   502  		var fValue reflect.Value
   503  		var fRet reflect.Value
   504  		if tag == nil {
   505  			// use field name
   506  			fValue = inValue.FieldByName(part)
   507  			if !fValue.IsValid() {
   508  				return fmt.Errorf("field name %v is not found in struct %v", part, inValue.Type().String())
   509  			}
   510  			fRet = ret.FieldByName(part)
   511  		} else {
   512  			// search tag that has key equal to part
   513  			found := false
   514  			for i := 0; i < inValue.NumField(); i++ {
   515  				f := inValue.Type().Field(i)
   516  				if key, ok := f.Tag.Lookup(*tag); ok {
   517  					if key == part {
   518  						fValue = inValue.Field(i)
   519  						fRet = ret.Field(i)
   520  						found = true
   521  						break
   522  					}
   523  				}
   524  			}
   525  			if !found {
   526  				return fmt.Errorf("Struct tag %v is not found with key %v", *tag, part)
   527  			}
   528  		}
   529  		// init Ret is zero and go down one more level
   530  		if fRet.IsZero() {
   531  			fRet.Set(reflect.New(fValue.Type()).Elem())
   532  		}
   533  		return prune(fValue, fRet, parts[1:], tag)
   534  	case reflect.Array, reflect.Slice:
   535  		// set all its elements
   536  		length := inValue.Len()
   537  		// init ret
   538  		if ret.IsZero() {
   539  			if inKind == reflect.Slice {
   540  				ret.Set(reflect.MakeSlice(inValue.Type(), length /*len*/, length /*cap*/))
   541  			} else { // array
   542  				ret.Set(reflect.New(inValue.Type()).Elem())
   543  			}
   544  		}
   545  		for j := 0; j < length; j++ {
   546  			if err := prune(inValue.Index(j), ret.Index(j), parts, tag); err != nil {
   547  				return err
   548  			}
   549  		}
   550  	default:
   551  		return fmt.Errorf("path %v cannot be looked up on kind of %v", strings.Join(parts, "."), inValue.Kind())
   552  	}
   553  
   554  	return nil
   555  }
   556
View as plain text