Resolve #40, init pivot table support

1 files changed, 434 insertions, 0 deletions

diff --git a/pivotTable.go b/pivotTable.go
new file mode 100644
index 0000000..881d774
--- /dev/null
+++ b/pivotTable.go
@@ -0,0 +1,434 @@
+// Copyright 2016 - 2019 The excelize Authors. All rights reserved. Use of
+// this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+//
+// Package excelize providing a set of functions that allow you to write to
+// and read from XLSX files. Support reads and writes XLSX file generated by
+// Microsoft Excel™ 2007 and later. Support save file without losing original
+// charts of XLSX. This library needs Go version 1.10 or later.
+
+package excelize
+
+import (
+	"encoding/xml"
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+// PivotTableOption directly maps the format settings of the pivot table.
+type PivotTableOption struct {
+	DataRange       string
+	PivotTableRange string
+	Rows            []string
+	Columns         []string
+	Data            []string
+	Page            []string
+}
+
+// AddPivotTable provides the method to add pivot table by given pivot table
+// options. For example, create a pivot table on the Sheet1!$G$2:$M$34 area
+// with the region Sheet1!$A$1:$E$31 as the data source, summarize by sum for
+// sales:
+//
+//    package main
+//
+//    import (
+//        "fmt"
+//        "math/rand"
+//
+//        "github.com/360EntSecGroup-Skylar/excelize"
+//    )
+//
+//    func main() {
+//        f := excelize.NewFile()
+//        // Create some data in a sheet
+//        month := []string{"Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"}
+//        year := []int{2017, 2018, 2019}
+//        types := []string{"Meat", "Dairy", "Beverages", "Produce"}
+//        region := []string{"East", "West", "North", "South"}
+//        f.SetSheetRow("Sheet1", "A1", &[]string{"Month", "Year", "Type", "Sales", "Region"})
+//        for i := 0; i < 30; i++ {
+//            f.SetCellValue("Sheet1", fmt.Sprintf("A%d", i+2), month[rand.Intn(12)])
+//            f.SetCellValue("Sheet1", fmt.Sprintf("B%d", i+2), year[rand.Intn(3)])
+//            f.SetCellValue("Sheet1", fmt.Sprintf("C%d", i+2), types[rand.Intn(4)])
+//            f.SetCellValue("Sheet1", fmt.Sprintf("D%d", i+2), rand.Intn(5000))
+//            f.SetCellValue("Sheet1", fmt.Sprintf("E%d", i+2), region[rand.Intn(4)])
+//        }
+//        err := f.AddPivotTable(&excelize.PivotTableOption{
+//            DataRange:       "Sheet1!$A$1:$E$31",
+//            PivotTableRange: "Sheet1!$G$2:$M$34",
+//            Rows:            []string{"Month", "Year"},
+//            Columns:         []string{"Type"},
+//            Data:            []string{"Sales"},
+//        })
+//        if err != nil {
+//            fmt.Println(err)
+//        }
+//        err = f.SaveAs("Book1.xlsx")
+//        if err != nil {
+//            fmt.Println(err)
+//        }
+//    }
+//
+func (f *File) AddPivotTable(opt *PivotTableOption) error {
+	// parameter validation
+	dataSheet, pivotTableSheetPath, err := f.parseFormatPivotTableSet(opt)
+	if err != nil {
+		return err
+	}
+
+	pivotTableID := f.countPivotTables() + 1
+	pivotCacheID := f.countPivotCache() + 1
+
+	sheetRelationshipsPivotTableXML := "../pivotTables/pivotTable" + strconv.Itoa(pivotTableID) + ".xml"
+	pivotTableXML := strings.Replace(sheetRelationshipsPivotTableXML, "..", "xl", -1)
+	pivotCacheXML := "xl/pivotCache/pivotCacheDefinition" + strconv.Itoa(pivotCacheID) + ".xml"
+	err = f.addPivotCache(pivotCacheID, pivotCacheXML, opt, dataSheet)
+	if err != nil {
+		return err
+	}
+
+	// workbook pivot cache
+	workBookPivotCacheRID := f.addRels("xl/_rels/workbook.xml.rels", SourceRelationshipPivotCache, fmt.Sprintf("pivotCache/pivotCacheDefinition%d.xml", pivotCacheID), "")
+	cacheID := f.addWorkbookPivotCache(workBookPivotCacheRID)
+
+	pivotCacheRels := "xl/pivotTables/_rels/pivotTable" + strconv.Itoa(pivotTableID) + ".xml.rels"
+	// rId not used
+	_ = f.addRels(pivotCacheRels, SourceRelationshipPivotCache, fmt.Sprintf("../pivotCache/pivotCacheDefinition%d.xml", pivotCacheID), "")
+	err = f.addPivotTable(cacheID, pivotTableID, pivotTableXML, opt)
+	if err != nil {
+		return err
+	}
+	pivotTableSheetRels := "xl/worksheets/_rels/" + strings.TrimPrefix(pivotTableSheetPath, "xl/worksheets/") + ".rels"
+	f.addRels(pivotTableSheetRels, SourceRelationshipPivotTable, sheetRelationshipsPivotTableXML, "")
+	f.addContentTypePart(pivotTableID, "pivotTable")
+	f.addContentTypePart(pivotCacheID, "pivotCache")
+
+	return nil
+}
+
+// parseFormatPivotTableSet provides a function to validate pivot table
+// properties.
+func (f *File) parseFormatPivotTableSet(opt *PivotTableOption) (*xlsxWorksheet, string, error) {
+	if opt == nil {
+		return nil, "", errors.New("parameter is required")
+	}
+	dataSheetName, _, err := f.adjustRange(opt.DataRange)
+	if err != nil {
+		return nil, "", fmt.Errorf("parameter 'DataRange' parsing error: %s", err.Error())
+	}
+	pivotTableSheetName, _, err := f.adjustRange(opt.PivotTableRange)
+	if err != nil {
+		return nil, "", fmt.Errorf("parameter 'PivotTableRange' parsing error: %s", err.Error())
+	}
+	dataSheet, err := f.workSheetReader(dataSheetName)
+	if err != nil {
+		return dataSheet, "", err
+	}
+	pivotTableSheetPath, ok := f.sheetMap[trimSheetName(pivotTableSheetName)]
+	if !ok {
+		return dataSheet, pivotTableSheetPath, fmt.Errorf("sheet %s is not exist", pivotTableSheetName)
+	}
+	return dataSheet, pivotTableSheetPath, err
+}
+
+// adjustRange adjust range, for example: adjust Sheet1!$E$31:$A$1 to Sheet1!$A$1:$E$31
+func (f *File) adjustRange(rangeStr string) (string, []int, error) {
+	if len(rangeStr) < 1 {
+		return "", []int{}, errors.New("parameter is required")
+	}
+	rng := strings.Split(rangeStr, "!")
+	if len(rng) != 2 {
+		return "", []int{}, errors.New("parameter is invalid")
+	}
+	trimRng := strings.Replace(rng[1], "$", "", -1)
+	coordinates, err := f.areaRefToCoordinates(trimRng)
+	if err != nil {
+		return rng[0], []int{}, err
+	}
+	x1, y1, x2, y2 := coordinates[0], coordinates[1], coordinates[2], coordinates[3]
+	if x1 == x2 && y1 == y2 {
+		return rng[0], []int{}, errors.New("parameter is invalid")
+	}
+
+	// Correct the coordinate area, such correct C1:B3 to B1:C3.
+	if x2 < x1 {
+		x1, x2 = x2, x1
+	}
+
+	if y2 < y1 {
+		y1, y2 = y2, y1
+	}
+	return rng[0], []int{x1, y1, x2, y2}, nil
+}
+
+func (f *File) getPivotFieldsOrder(dataRange string) ([]string, error) {
+	order := []string{}
+	// data range has been checked
+	dataSheet, coordinates, err := f.adjustRange(dataRange)
+	if err != nil {
+		return order, fmt.Errorf("parameter 'DataRange' parsing error: %s", err.Error())
+	}
+	for col := coordinates[0]; col <= coordinates[2]; col++ {
+		coordinate, _ := CoordinatesToCellName(col, coordinates[1])
+		name, err := f.GetCellValue(dataSheet, coordinate)
+		if err != nil {
+			return order, err
+		}
+		order = append(order, name)
+	}
+	return order, nil
+}
+
+// addPivotCache provides a function to create a pivot cache by given properties.
+func (f *File) addPivotCache(pivotCacheID int, pivotCacheXML string, opt *PivotTableOption, ws *xlsxWorksheet) error {
+	// validate data range
+	dataSheet, coordinates, err := f.adjustRange(opt.DataRange)
+	if err != nil {
+		return fmt.Errorf("parameter 'DataRange' parsing error: %s", err.Error())
+	}
+	order, err := f.getPivotFieldsOrder(opt.DataRange)
+	if err != nil {
+		return err
+	}
+	hcell, _ := CoordinatesToCellName(coordinates[0], coordinates[1])
+	vcell, _ := CoordinatesToCellName(coordinates[2], coordinates[3])
+	pc := xlsxPivotCacheDefinition{
+		SaveData:      false,
+		RefreshOnLoad: true,
+		CacheSource: &xlsxCacheSource{
+			Type: "worksheet",
+			WorksheetSource: &xlsxWorksheetSource{
+				Ref:   hcell + ":" + vcell,
+				Sheet: dataSheet,
+			},
+		},
+		CacheFields: &xlsxCacheFields{},
+	}
+	for _, name := range order {
+		pc.CacheFields.CacheField = append(pc.CacheFields.CacheField, &xlsxCacheField{
+			Name: name,
+			SharedItems: &xlsxSharedItems{
+				Count: 0,
+			},
+		})
+	}
+	pc.CacheFields.Count = len(pc.CacheFields.CacheField)
+	pivotCache, err := xml.Marshal(pc)
+	f.saveFileList(pivotCacheXML, pivotCache)
+	return err
+}
+
+// addPivotTable provides a function to create a pivot table by given pivot
+// table ID and properties.
+func (f *File) addPivotTable(cacheID, pivotTableID int, pivotTableXML string, opt *PivotTableOption) error {
+	// validate pivot table range
+	_, coordinates, err := f.adjustRange(opt.PivotTableRange)
+	if err != nil {
+		return fmt.Errorf("parameter 'PivotTableRange' parsing error: %s", err.Error())
+	}
+
+	hcell, _ := CoordinatesToCellName(coordinates[0], coordinates[1])
+	vcell, _ := CoordinatesToCellName(coordinates[2], coordinates[3])
+
+	pt := xlsxPivotTableDefinition{
+		Name:        fmt.Sprintf("Pivot Table%d", pivotTableID),
+		CacheID:     cacheID,
+		DataCaption: "Values",
+		Location: &xlsxLocation{
+			Ref:            hcell + ":" + vcell,
+			FirstDataCol:   1,
+			FirstDataRow:   1,
+			FirstHeaderRow: 1,
+		},
+		PivotFields: &xlsxPivotFields{},
+		RowFields:   &xlsxRowFields{},
+		RowItems: &xlsxRowItems{
+			Count: 1,
+			I: []*xlsxI{
+				{
+					[]*xlsxX{{}, {}},
+				},
+			},
+		},
+		ColFields:  &xlsxColFields{},
+		DataFields: &xlsxDataFields{},
+		PivotTableStyleInfo: &xlsxPivotTableStyleInfo{
+			Name:           "PivotStyleLight16",
+			ShowRowHeaders: true,
+			ShowColHeaders: true,
+			ShowLastColumn: true,
+		},
+	}
+
+	// pivot fields
+	err = f.addPivotFields(&pt, opt)
+	if err != nil {
+		return err
+	}
+
+	// count pivot fields
+	pt.PivotFields.Count = len(pt.PivotFields.PivotField)
+
+	// row fields
+	rowFieldsIndex, err := f.getPivotFieldsIndex(opt.Rows, opt)
+	if err != nil {
+		return err
+	}
+	for _, filedIdx := range rowFieldsIndex {
+		pt.RowFields.Field = append(pt.RowFields.Field, &xlsxField{
+			X: filedIdx,
+		})
+	}
+
+	// count row fields
+	pt.RowFields.Count = len(pt.RowFields.Field)
+
+	// col fields
+	colFieldsIndex, err := f.getPivotFieldsIndex(opt.Columns, opt)
+	if err != nil {
+		return err
+	}
+	for _, filedIdx := range colFieldsIndex {
+		pt.ColFields.Field = append(pt.ColFields.Field, &xlsxField{
+			X: filedIdx,
+		})
+	}
+
+	// count col fields
+	pt.ColFields.Count = len(pt.ColFields.Field)
+
+	// data fields
+	dataFieldsIndex, err := f.getPivotFieldsIndex(opt.Data, opt)
+	if err != nil {
+		return err
+	}
+	for _, dataField := range dataFieldsIndex {
+		pt.DataFields.DataField = append(pt.DataFields.DataField, &xlsxDataField{
+			Fld: dataField,
+		})
+	}
+
+	// count data fields
+	pt.DataFields.Count = len(pt.DataFields.DataField)
+
+	pivotTable, err := xml.Marshal(pt)
+	f.saveFileList(pivotTableXML, pivotTable)
+	return err
+}
+
+// inStrSlice provides a method to check if an element is present in an array,
+// and return the index of its location, otherwise return -1.
+func inStrSlice(a []string, x string) int {
+	for idx, n := range a {
+		if x == n {
+			return idx
+		}
+	}
+	return -1
+}
+
+// addPivotFields create pivot fields based on the column order of the first
+// row in the data region by given pivot table definition and option.
+func (f *File) addPivotFields(pt *xlsxPivotTableDefinition, opt *PivotTableOption) error {
+	order, err := f.getPivotFieldsOrder(opt.DataRange)
+	if err != nil {
+		return err
+	}
+	for _, name := range order {
+		if inStrSlice(opt.Rows, name) != -1 {
+			pt.PivotFields.PivotField = append(pt.PivotFields.PivotField, &xlsxPivotField{
+				Axis: "axisRow",
+				Items: &xlsxItems{
+					Count: 1,
+					Item: []*xlsxItem{
+						{T: "default"},
+					},
+				},
+			})
+			continue
+		}
+		if inStrSlice(opt.Columns, name) != -1 {
+			pt.PivotFields.PivotField = append(pt.PivotFields.PivotField, &xlsxPivotField{
+				Axis: "axisCol",
+				Items: &xlsxItems{
+					Count: 1,
+					Item: []*xlsxItem{
+						{T: "default"},
+					},
+				},
+			})
+			continue
+		}
+		if inStrSlice(opt.Data, name) != -1 {
+			pt.PivotFields.PivotField = append(pt.PivotFields.PivotField, &xlsxPivotField{
+				DataField: true,
+			})
+			continue
+		}
+		pt.PivotFields.PivotField = append(pt.PivotFields.PivotField, &xlsxPivotField{})
+	}
+	return err
+}
+
+// countPivotTables provides a function to get drawing files count storage in
+// the folder xl/pivotTables.
+func (f *File) countPivotTables() int {
+	count := 0
+	for k := range f.XLSX {
+		if strings.Contains(k, "xl/pivotTables/pivotTable") {
+			count++
+		}
+	}
+	return count
+}
+
+// countPivotCache provides a function to get drawing files count storage in
+// the folder xl/pivotCache.
+func (f *File) countPivotCache() int {
+	count := 0
+	for k := range f.XLSX {
+		if strings.Contains(k, "xl/pivotCache/pivotCacheDefinition") {
+			count++
+		}
+	}
+	return count
+}
+
+// getPivotFieldsIndex convert the column of the first row in the data region
+// to a sequential index by given fields and pivot option.
+func (f *File) getPivotFieldsIndex(fields []string, opt *PivotTableOption) ([]int, error) {
+	pivotFieldsIndex := []int{}
+	orders, err := f.getPivotFieldsOrder(opt.DataRange)
+	if err != nil {
+		return pivotFieldsIndex, err
+	}
+	for _, field := range fields {
+		if pos := inStrSlice(orders, field); pos != -1 {
+			pivotFieldsIndex = append(pivotFieldsIndex, pos)
+		}
+	}
+	return pivotFieldsIndex, nil
+}
+
+// addWorkbookPivotCache add the association ID of the pivot cache in xl/workbook.xml.
+func (f *File) addWorkbookPivotCache(RID int) int {
+	wb := f.workbookReader()
+	if wb.PivotCaches == nil {
+		wb.PivotCaches = &xlsxPivotCaches{}
+	}
+	cacheID := 1
+	for _, pivotCache := range wb.PivotCaches.PivotCache {
+		if pivotCache.CacheID > cacheID {
+			cacheID = pivotCache.CacheID
+		}
+	}
+	cacheID++
+	wb.PivotCaches.PivotCache = append(wb.PivotCaches.PivotCache, xlsxPivotCache{
+		CacheID: cacheID,
+		RID:     fmt.Sprintf("rId%d", RID),
+	})
+	return cacheID
+}