Add custom ZIP writer support for WriteToBuffer function

excelize.go

@@ -17,6 +17,7 @@ import (
 	"bytes"
 	"encoding/xml"
 	"io"
+	"io/fs"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -41,7 +42,7 @@ type File struct {
 	tempFiles        sync.Map
 	xmlAttr          sync.Map
 	CalcChain        *xlsxCalcChain
-	CharsetReader    charsetTranscoderFn
+	CharsetReader    func(charset string, input io.Reader) (rdr io.Reader, err error)
 	Comments         map[string]*xlsxComments
 	ContentTypes     *xlsxTypes
 	DecodeVMLDrawing map[string]*decodeVmlDrawing
@@ -58,11 +59,17 @@ type File struct {
 	VMLDrawing       map[string]*vmlDrawing
 	VolatileDeps     *xlsxVolTypes
 	WorkBook         *xlsxWorkbook
+	ZipWriter        func(io.Writer) ZipWriter
 }
 
-// charsetTranscoderFn set user-defined codepage transcoder function for open
-// the spreadsheet from non-UTF-8 encoding.
-type charsetTranscoderFn func(charset string, input io.Reader) (rdr io.Reader, err error)
+// ZipWriter defines an interface for writing files to a ZIP archive. It
+// provides methods to create new files within the archive, add files from a
+// filesystem, and close the archive when writing is complete.
+type ZipWriter interface {
+	Create(name string) (io.Writer, error)
+	AddFS(fsys fs.FS) error
+	Close() error
+}
 
 // Options define the options for opening and reading the spreadsheet.
 //
@@ -153,6 +160,7 @@ func newFile() *File {
 		VMLDrawing:       make(map[string]*vmlDrawing),
 		Relationships:    sync.Map{},
 		CharsetReader:    charset.NewReaderLabel,
+		ZipWriter:        func(w io.Writer) ZipWriter { return zip.NewWriter(w) },
 	}
 }
 
@@ -232,9 +240,15 @@ func (f *File) getOptions(opts ...Options) *Options {
 	return options
 }
 
-// CharsetTranscoder Set user defined codepage transcoder function for open
+// CharsetTranscoder set user defined codepage transcoder function for open
 // workbook from non UTF-8 encoding.
-func (f *File) CharsetTranscoder(fn charsetTranscoderFn) *File { f.CharsetReader = fn; return f }
+func (f *File) CharsetTranscoder(fn func(charset string, input io.Reader) (rdr io.Reader, err error)) *File {
+	f.CharsetReader = fn
+	return f
+}
+
+// SetZipWriter set user defined zip writer function for saving the workbook.
+func (f *File) SetZipWriter(fn func(io.Writer) ZipWriter) *File { f.ZipWriter = fn; return f }
 
 // Creates new XML decoder with charset reader.
 func (f *File) xmlNewDecoder(rdr io.Reader) (ret *xml.Decoder) {

excelize_test.go

@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"golang.org/x/net/html/charset"
 )
 
 func TestOpenFile(t *testing.T) {
@@ -244,7 +245,7 @@ func TestSaveAsWrongPath(t *testing.T) {
 
 func TestCharsetTranscoder(t *testing.T) {
 	f := NewFile()
-	f.CharsetTranscoder(*new(charsetTranscoderFn))
+	f.CharsetTranscoder(charset.NewReaderLabel)
 }
 
 func TestOpenReader(t *testing.T) {
@@ -355,7 +356,7 @@ func TestOpenReader(t *testing.T) {
 
 func TestBrokenFile(t *testing.T) {
 	// Test write file with broken file struct
-	f := File{}
+	f := File{ZipWriter: func(w io.Writer) ZipWriter { return zip.NewWriter(w) }}
 
 	t.Run("SaveWithoutName", func(t *testing.T) {
 		assert.EqualError(t, f.Save(), "no path defined for file, consider File.WriteTo or File.Write")

file.go

@@ -12,7 +12,6 @@
 package excelize
 
 import (
-	"archive/zip"
 	"bytes"
 	"encoding/binary"
 	"encoding/xml"
@@ -137,7 +136,7 @@ func (f *File) WriteTo(w io.Writer, opts ...Options) (int64, error) {
 // and it allocates space in memory. Be careful when the file size is large.
 func (f *File) WriteToBuffer() (*bytes.Buffer, error) {
 	buf := new(bytes.Buffer)
-	zw := zip.NewWriter(buf)
+	zw := f.ZipWriter(buf)
 
 	if err := f.writeToZip(zw); err != nil {
 		_ = zw.Close()
@@ -158,8 +157,8 @@ func (f *File) WriteToBuffer() (*bytes.Buffer, error) {
 	return buf, nil
 }
 
-// writeToZip provides a function to write to zip.Writer
-func (f *File) writeToZip(zw *zip.Writer) error {
+// writeToZip provides a function to write to ZipWriter.
+func (f *File) writeToZip(zw ZipWriter) error {
 	f.calcChainWriter()
 	f.commentsWriter()
 	f.contentTypesWriter()

file_test.go

@@ -1,9 +1,11 @@
 package excelize
 
 import (
+	"archive/zip"
 	"bufio"
 	"bytes"
 	"encoding/binary"
+	"io"
 	"math"
 	"os"
 	"path/filepath"
@@ -43,6 +45,7 @@ func TestWriteTo(t *testing.T) {
 	// Test WriteToBuffer err
 	{
 		f, buf := File{Pkg: sync.Map{}}, bytes.Buffer{}
+		f.SetZipWriter(func(w io.Writer) ZipWriter { return zip.NewWriter(w) })
 		f.Pkg.Store("/d/", []byte("s"))
 		_, err := f.WriteTo(bufio.NewWriter(&buf))
 		assert.EqualError(t, err, "zip: write to directory")
@@ -51,6 +54,7 @@ func TestWriteTo(t *testing.T) {
 	// Test file path overflow
 	{
 		f, buf := File{Pkg: sync.Map{}}, bytes.Buffer{}
+		f.SetZipWriter(func(w io.Writer) ZipWriter { return zip.NewWriter(w) })
 		const maxUint16 = 1<<16 - 1
 		f.Pkg.Store(strings.Repeat("s", maxUint16+1), nil)
 		_, err := f.WriteTo(bufio.NewWriter(&buf))
@@ -59,6 +63,7 @@ func TestWriteTo(t *testing.T) {
 	// Test StreamsWriter err
 	{
 		f, buf := File{Pkg: sync.Map{}}, bytes.Buffer{}
+		f.SetZipWriter(func(w io.Writer) ZipWriter { return zip.NewWriter(w) })
 		f.Pkg.Store("s", nil)
 		f.streams = make(map[string]*StreamWriter)
 		file, _ := os.Open("123")
@@ -69,6 +74,7 @@ func TestWriteTo(t *testing.T) {
 	// Test write with temporary file
 	{
 		f, buf := File{tempFiles: sync.Map{}}, bytes.Buffer{}
+		f.SetZipWriter(func(w io.Writer) ZipWriter { return zip.NewWriter(w) })
 		const maxUint16 = 1<<16 - 1
 		f.tempFiles.Store("s", "")
 		f.tempFiles.Store(strings.Repeat("s", maxUint16+1), "")
@@ -78,6 +84,7 @@ func TestWriteTo(t *testing.T) {
 	// Test write with unsupported workbook file format
 	{
 		f, buf := File{Pkg: sync.Map{}}, bytes.Buffer{}
+		f.SetZipWriter(func(w io.Writer) ZipWriter { return zip.NewWriter(w) })
 		f.Pkg.Store("/d", []byte("s"))
 		f.Path = "Book1.xls"
 		_, err := f.WriteTo(bufio.NewWriter(&buf))