go语言 array_column 方法（反射）

php 的 array_column 太太太方便了，实现了一款 go 语言版本供借鉴

相关代码放到了github：https://github.com/entertainment-venue/poketto

 

功能1：提取 []struct 中 column 列。desk中存储为 slice

功能2：提取 []struct 中的 index 列作为 key，column列作为值。 desk 中存储为map

 

参数说明 :

@param desk [slice|map] 指针类型，方法最终的存储位置

@param input []struct，待转换的结构体切片 

@param columnKey string

@param indexKey string

 

go语言代码实现：

package struct_column

import (
	"errors"
	"fmt"
	"reflect"
)

/**
 * len(indexKey) > 0 将切片类型的结构体，转成 map 输出，输出结果存在了 desk 中，也就是你的第一个参数
 * len(indexKey) == 0 将切片类型的结构体，转成 slice 输出，输出结果存在了 desk 中，也就是你的第一个参数
 *
 * @demo 假如有下述结构体，
 * type User struct {
 * 	ID   int
 * 	NAME string
 * }
 * 输入：
 * @params desk &map[int64]User   desk是个指针呦！！！
 * @params input []User{User{ID:1, NAME:"zwk"}, User{ID:2, NAME:"zzz"}}
 * @params indexKey 键名 "ID"
 * @params columnKey 列名 空字符串代表返回整个结构体，反之返回结构体中的某一列
 *
 * 输出：
 * err 错误信息
 * 入参 desk 已经被赋值：map[int]User{1:User{ID:1, NAME:"zwk"}, 2:User{ID:2, NAME:"zzz"}}
 *
 * 输入：
 * @params desk &[]int   desk是个指针呦！！！
 * @params input []User{User{ID:1, NAME:"zwk"}, User{ID:2, NAME:"zzz"}}
 * @params indexKey 键名 ""
 * @params columnKey 列名
 *
 * 输出：
 * err 错误信息
 * 入参 desk 已经被赋值：[]int{1, 2}
 */
func StructColumn(desk, input interface{}, columnKey, indexKey string) (err error) {
	structIndexColumn := func(desk, input interface{}, columnKey, indexKey string) (err error) {
		findStructValByIndexKey := func(curVal reflect.Value, elemType reflect.Type, indexKey, columnKey string) (indexVal, columnVal reflect.Value, err error) {
			indexExist := false
			columnExist := false
			for i := 0; i < elemType.NumField(); i++ {
				curField := curVal.Field(i)
				if elemType.Field(i).Name == indexKey {
					switch curField.Kind() {
					case reflect.String, reflect.Int64, reflect.Int32, reflect.Int16, reflect.Int8, reflect.Int, reflect.Float64, reflect.Float32:
						indexExist = true
						indexVal = curField
					default:
						return indexVal, columnVal, errors.New("indexKey must be int float or string")
					}
				}
				if elemType.Field(i).Name == columnKey {
					columnExist = true
					columnVal = curField
					continue
				}
			}
			if !indexExist {
				return indexVal, columnVal, errors.New(fmt.Sprintf("indexKey %s not found in %s's field", indexKey, elemType))
			}
			if len(columnKey) > 0 && !columnExist {
				return indexVal, columnVal, errors.New(fmt.Sprintf("columnKey %s not found in %s's field", columnKey, elemType))
			}
			return
		}

		deskValue := reflect.ValueOf(desk)
		if deskValue.Elem().Kind() != reflect.Map {
			return errors.New("desk must be map")
		}
		deskElem := deskValue.Type().Elem()
		if len(columnKey) == 0 && deskElem.Elem().Kind() != reflect.Struct {
			return errors.New(fmt.Sprintf("desk's elem expect struct, got %s", deskElem.Elem().Kind()))
		}

		rv := reflect.ValueOf(input)
		rt := reflect.TypeOf(input)
		elemType := rt.Elem()

		var indexVal, columnVal reflect.Value
		direct := reflect.Indirect(deskValue)
		mapReflect := reflect.MakeMap(deskElem)
		deskKey := deskValue.Type().Elem().Key()

		for i := 0; i < rv.Len(); i++ {
			curVal := rv.Index(i)
			indexVal, columnVal, err = findStructValByIndexKey(curVal, elemType, indexKey, columnKey)
			if err != nil {
				return
			}
			if deskKey.Kind() != indexVal.Kind() {
				return errors.New(fmt.Sprintf("cant't convert %s to %s, your map'key must be %s", indexVal.Kind(), deskKey.Kind(), indexVal.Kind()))
			}
			if len(columnKey) == 0 {
				mapReflect.SetMapIndex(indexVal, curVal)
				direct.Set(mapReflect)
			} else {
				if deskElem.Elem().Kind() != columnVal.Kind() {
					return errors.New(fmt.Sprintf("your map must be map[%s]%s", indexVal.Kind(), columnVal.Kind()))
				}
				mapReflect.SetMapIndex(indexVal, columnVal)
				direct.Set(mapReflect)
			}
		}
		return
	}

	structColumn := func(desk, input interface{}, columnKey string) (err error) {
		findStructValByColumnKey := func(curVal reflect.Value, elemType reflect.Type, columnKey string) (columnVal reflect.Value, err error) {
			columnExist := false
			for i := 0; i < elemType.NumField(); i++ {
				curField := curVal.Field(i)
				if elemType.Field(i).Name == columnKey {
					columnExist = true
					columnVal = curField
					continue
				}
			}
			if !columnExist {
				return columnVal, errors.New(fmt.Sprintf("columnKey %s not found in %s's field", columnKey, elemType))
			}
			return
		}

		if len(columnKey) == 0 {
			return errors.New("columnKey cannot not be empty")
		}

		deskElemType := reflect.TypeOf(desk).Elem()
		if deskElemType.Kind() != reflect.Slice {
			return errors.New("desk must be slice")
		}

		rv := reflect.ValueOf(input)
		rt := reflect.TypeOf(input)

		var columnVal reflect.Value
		deskValue := reflect.ValueOf(desk)
		direct := reflect.Indirect(deskValue)

		for i := 0; i < rv.Len(); i++ {
			columnVal, err = findStructValByColumnKey(rv.Index(i), rt.Elem(), columnKey)
			if err != nil {
				return
			}
			if deskElemType.Elem().Kind() != columnVal.Kind() {
				return errors.New(fmt.Sprintf("your slice must be []%s", columnVal.Kind()))
			}

			direct.Set(reflect.Append(direct, columnVal))
		}
		return
	}

	deskValue := reflect.ValueOf(desk)
	if deskValue.Kind() != reflect.Ptr {
		return errors.New("desk must be ptr")
	}

	rv := reflect.ValueOf(input)
	if rv.Kind() != reflect.Slice && rv.Kind() != reflect.Array {
		return errors.New("input must be map slice or array")
	}

	rt := reflect.TypeOf(input)
	if rt.Elem().Kind() != reflect.Struct {
		return errors.New("input's elem must be struct")
	}

	if len(indexKey) > 0 {
		return structIndexColumn(desk, input, columnKey, indexKey)
	}
	return structColumn(desk, input, columnKey)
}

 

 

测试

package struct_column

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

type User struct {
	Id   int
	Name string
}

func TestStructColumnForMap(t *testing.T) {
	u1 := User{
		Id:   1,
		Name: "uzi",
	}
	u2 := User{
		Id:   2,
		Name: "faker",
	}
	users := []User{u1, u2}

	var desk map[int]User

	if err := StructColumn(&desk, users, "", "Id"); err != nil {
		assert.Fail(t, err.Error())
		return
	}
	assert.Equal(t, len(desk), 2)
	assert.Equal(t, desk[1].Name, u1.Name)
	assert.Equal(t, desk[2].Name, u2.Name)

	var desk2 map[int]string

	if err := StructColumn(&desk2, users, "Name", "Id"); err != nil {
		assert.Fail(t, err.Error())
		return
	}
	assert.Equal(t, len(desk2), 2)
	assert.Equal(t, desk2[1], u1.Name)
	assert.Equal(t, desk2[2], u2.Name)
}

func TestStructColumnForSlice(t *testing.T) {
	u1 := User{
		Id:   1,
		Name: "uzi",
	}
	u2 := User{
		Id:   2,
		Name: "faker",
	}
	users := []User{u1, u2}

	var desk []string

	if err := StructColumn(&desk, users, "Name", ""); err != nil {
		assert.Fail(t, err.Error())
		return
	}
	assert.Equal(t, len(desk), 2)
	assert.Equal(t, desk[0], "uzi")
	assert.Equal(t, desk[1], "faker")
}