Resolve #217, new function add VBA project supported.

4 files changed, 89 insertions, 0 deletions

diff --git a/excelize.go b/excelize.go
index f636a84..c7eff10 100644
--- a/excelize.go
+++ b/excelize.go
@@ -19,7 +19,9 @@ import (
 	"io"
 	"io/ioutil"
 	"os"
+	"path"
 	"strconv"
+	"strings"
 )
 
 // File define a populated XLSX file struct.
@@ -226,3 +228,78 @@ func (f *File) UpdateLinkedValue() error {
 	}
 	return nil
 }
+
+// AddVBAProject provides the method to add vbaProject.bin file which contains
+// functions and/or macros. The file extension should be .xlsm. For example:
+//
+//    err := f.SetSheetPrOptions("Sheet1", excelize.CodeName("Sheet1"))
+//    if err != nil {
+//        fmt.Println(err)
+//    }
+//    err = f.AddVBAProject("vbaProject.bin")
+//    if err != nil {
+//        fmt.Println(err)
+//    }
+//    err = f.SaveAs("macros.xlsm")
+//    if err != nil {
+//        fmt.Println(err)
+//    }
+//
+func (f *File) AddVBAProject(bin string) error {
+	var err error
+	// Check vbaProject.bin exists first.
+	if _, err = os.Stat(bin); os.IsNotExist(err) {
+		return err
+	}
+	if path.Ext(bin) != ".bin" {
+		return errors.New("unsupported VBA project extension")
+	}
+	f.setContentTypePartVBAProjectExtensions()
+	wb := f.workbookRelsReader()
+	var rID int
+	var ok bool
+	for _, rel := range wb.Relationships {
+		if rel.Target == "vbaProject.bin" && rel.Type == SourceRelationshipVBAProject {
+			ok = true
+			continue
+		}
+		t, _ := strconv.Atoi(strings.TrimPrefix(rel.ID, "rId"))
+		if t > rID {
+			rID = t
+		}
+	}
+	rID++
+	if !ok {
+		wb.Relationships = append(wb.Relationships, xlsxWorkbookRelation{
+			ID:     "rId" + strconv.Itoa(rID),
+			Target: "vbaProject.bin",
+			Type:   SourceRelationshipVBAProject,
+		})
+	}
+	file, _ := ioutil.ReadFile(bin)
+	f.XLSX["xl/vbaProject.bin"] = file
+	return err
+}
+
+// setContentTypePartVBAProjectExtensions provides a function to set the
+// content type for relationship parts and the main document part.
+func (f *File) setContentTypePartVBAProjectExtensions() {
+	var ok bool
+	content := f.contentTypesReader()
+	for _, v := range content.Defaults {
+		if v.Extension == "bin" {
+			ok = true
+		}
+	}
+	for idx, o := range content.Overrides {
+		if o.PartName == "/xl/workbook.xml" {
+			content.Overrides[idx].ContentType = "application/vnd.ms-excel.sheet.macroEnabled.main+xml"
+		}
+	}
+	if !ok {
+		content.Defaults = append(content.Defaults, xlsxDefault{
+			Extension:   "bin",
+			ContentType: "application/vnd.ms-office.vbaProject",
+		})
+	}
+}
diff --git a/excelize_test.go b/excelize_test.go
index c4a06a5..79010b1 100644
--- a/excelize_test.go
+++ b/excelize_test.go
@@ -1078,6 +1078,17 @@ func TestSetDefaultTimeStyle(t *testing.T) {
 	assert.EqualError(t, f.setDefaultTimeStyle("SheetN", "", 0), "sheet SheetN is not exist")
 }
 
+func TestAddVBAProject(t *testing.T) {
+	f := NewFile()
+	assert.NoError(t, f.SetSheetPrOptions("Sheet1", CodeName("Sheet1")))
+	assert.EqualError(t, f.AddVBAProject("macros.bin"), "stat macros.bin: no such file or directory")
+	assert.EqualError(t, f.AddVBAProject(filepath.Join("test", "Book1.xlsx")), "unsupported VBA project extension")
+	assert.NoError(t, f.AddVBAProject(filepath.Join("test", "vbaProject.bin")))
+	// Test add VBA project twice.
+	assert.NoError(t, f.AddVBAProject(filepath.Join("test", "vbaProject.bin")))
+	assert.NoError(t, f.SaveAs(filepath.Join("test", "TestAddVBAProject.xlsm")))
+}
+
 func prepareTestBook1() (*File, error) {
 	f, err := OpenFile(filepath.Join("test", "Book1.xlsx"))
 	if err != nil {
diff --git a/test/vbaProject.bin b/test/vbaProject.bin
new file mode 100755
index 0000000..fc15dca
--- /dev/null
+++ b/test/vbaProject.bin
diff --git a/xmlDrawing.go b/xmlDrawing.go
index bb468bc..1201cc8 100644
--- a/xmlDrawing.go
+++ b/xmlDrawing.go
@@ -22,6 +22,7 @@ const (
 	SourceRelationshipDrawingVML         = "http://schemas.openxmlformats.org/officeDocument/2006/relationships/vmlDrawing"
 	SourceRelationshipHyperLink          = "http://schemas.openxmlformats.org/officeDocument/2006/relationships/hyperlink"
 	SourceRelationshipWorkSheet          = "http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet"
+	SourceRelationshipVBAProject         = "http://schemas.microsoft.com/office/2006/relationships/vbaProject"
 	SourceRelationshipChart201506        = "http://schemas.microsoft.com/office/drawing/2015/06/chart"
 	SourceRelationshipChart20070802      = "http://schemas.microsoft.com/office/drawing/2007/8/2/chart"
 	SourceRelationshipChart2014          = "http://schemas.microsoft.com/office/drawing/2014/chart"