dataset.go 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. package xorm
  2. import (
  3. tablib "github.com/agrison/go-tablib"
  4. )
  5. // NewDataset creates a new Dataset.
  6. func NewDataset(headers []string) *tablib.Dataset {
  7. return tablib.NewDataset(headers)
  8. }
  9. // NewDatasetWithData creates a new Dataset.
  10. func NewDatasetWithData(headers []string, data interface{}) (*tablib.Dataset, error) {
  11. if data == nil {
  12. return tablib.NewDatasetWithData(headers, nil), nil
  13. }
  14. n := len(headers)
  15. switch data.(type) {
  16. case [][]interface{}:
  17. return tablib.NewDatasetWithData(headers, data.([][]interface{})), nil
  18. case []map[string]interface{}:
  19. dataSlice := data.([]map[string]interface{})
  20. if len(dataSlice) > 0 {
  21. if len(dataSlice[0]) == 0 {
  22. return tablib.NewDatasetWithData(headers, make([][]interface{}, len(dataSlice))), nil
  23. } else {
  24. if n != len(dataSlice[0]) {
  25. return nil, ErrParamsType
  26. }
  27. mapHeaders := make(map[string]int, n)
  28. for i := 0; i < n; i++ {
  29. mapHeaders[headers[i]] = i
  30. }
  31. for k, _ := range dataSlice[0] {
  32. if _, ok := mapHeaders[k]; !ok {
  33. return nil, ErrParamsType
  34. }
  35. }
  36. d := tablib.NewDataset(headers)
  37. var row []interface{}
  38. for i, _ := range dataSlice {
  39. row = nil
  40. for j := 0; j < n; j++ {
  41. row = append(row, dataSlice[i][headers[j]])
  42. }
  43. d.Append(row)
  44. }
  45. return d, nil
  46. }
  47. } else {
  48. return tablib.NewDatasetWithData(headers, nil), nil
  49. }
  50. return nil, ErrParamsType
  51. default:
  52. return nil, ErrParamsType
  53. }
  54. }