feat: stream write to zip directly (#863)

This commit is contained in:
Zitao 2021-06-22 14:06:08 +08:00 committed by GitHub
parent 2cfcf9eb5f
commit 24967a5c25
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 102 additions and 59 deletions

118
file.go
View File

@ -87,62 +87,26 @@ func (f *File) Write(w io.Writer) error {
// WriteTo implements io.WriterTo to write the file.
func (f *File) WriteTo(w io.Writer) (int64, error) {
buf, err := f.WriteToBuffer()
if err != nil {
if f.options != nil && f.options.Password != "" {
buf, err := f.WriteToBuffer()
if err != nil {
return 0, err
}
return buf.WriteTo(w)
}
if err := f.writeDirectToWriter(w); err != nil {
return 0, err
}
return buf.WriteTo(w)
return 0, nil
}
// WriteToBuffer provides a function to get bytes.Buffer from the saved file.
// WriteToBuffer provides a function to get bytes.Buffer from the saved file. And it allocate 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)
f.calcChainWriter()
f.commentsWriter()
f.contentTypesWriter()
f.drawingsWriter()
f.vmlDrawingWriter()
f.workBookWriter()
f.workSheetWriter()
f.relsWriter()
f.sharedStringsWriter()
f.styleSheetWriter()
for path, stream := range f.streams {
fi, err := zw.Create(path)
if err != nil {
zw.Close()
return buf, err
}
var from io.Reader
from, err = stream.rawData.Reader()
if err != nil {
stream.rawData.Close()
return buf, err
}
_, err = io.Copy(fi, from)
if err != nil {
zw.Close()
return buf, err
}
stream.rawData.Close()
}
for path, content := range f.XLSX {
if _, ok := f.streams[path]; ok {
continue
}
fi, err := zw.Create(path)
if err != nil {
zw.Close()
return buf, err
}
_, err = fi.Write(content)
if err != nil {
zw.Close()
return buf, err
}
if err := f.writeToZip(zw); err != nil {
return buf, zw.Close()
}
if f.options != nil && f.options.Password != "" {
@ -159,3 +123,61 @@ func (f *File) WriteToBuffer() (*bytes.Buffer, error) {
}
return buf, zw.Close()
}
// writeDirectToWriter provides a function to write to io.Writer.
func (f *File) writeDirectToWriter(w io.Writer) error {
zw := zip.NewWriter(w)
if err := f.writeToZip(zw); err != nil {
zw.Close()
return err
}
return zw.Close()
}
// writeToZip provides a function to write to zip.Writer
func (f *File) writeToZip(zw *zip.Writer) error {
f.calcChainWriter()
f.commentsWriter()
f.contentTypesWriter()
f.drawingsWriter()
f.vmlDrawingWriter()
f.workBookWriter()
f.workSheetWriter()
f.relsWriter()
f.sharedStringsWriter()
f.styleSheetWriter()
for path, stream := range f.streams {
fi, err := zw.Create(path)
if err != nil {
return err
}
var from io.Reader
from, err = stream.rawData.Reader()
if err != nil {
stream.rawData.Close()
return err
}
_, err = io.Copy(fi, from)
if err != nil {
return err
}
stream.rawData.Close()
}
for path, content := range f.XLSX {
if _, ok := f.streams[path]; ok {
continue
}
fi, err := zw.Create(path)
if err != nil {
return err
}
_, err = fi.Write(content)
if err != nil {
return err
}
}
return nil
}

View File

@ -3,6 +3,7 @@ package excelize
import (
"bufio"
"bytes"
"os"
"strings"
"testing"
@ -33,16 +34,36 @@ func BenchmarkWrite(b *testing.B) {
}
func TestWriteTo(t *testing.T) {
f := File{}
buf := bytes.Buffer{}
f.XLSX = make(map[string][]byte)
f.XLSX["/d/"] = []byte("s")
_, err := f.WriteTo(bufio.NewWriter(&buf))
assert.EqualError(t, err, "zip: write to directory")
delete(f.XLSX, "/d/")
// Test WriteToBuffer err
{
f := File{}
buf := bytes.Buffer{}
f.XLSX = make(map[string][]byte)
f.XLSX["/d/"] = []byte("s")
_, err := f.WriteTo(bufio.NewWriter(&buf))
assert.EqualError(t, err, "zip: write to directory")
delete(f.XLSX, "/d/")
}
// Test file path overflow
const maxUint16 = 1<<16 - 1
f.XLSX[strings.Repeat("s", maxUint16+1)] = nil
_, err = f.WriteTo(bufio.NewWriter(&buf))
assert.EqualError(t, err, "zip: FileHeader.Name too long")
{
f := File{}
buf := bytes.Buffer{}
f.XLSX = make(map[string][]byte)
const maxUint16 = 1<<16 - 1
f.XLSX[strings.Repeat("s", maxUint16+1)] = nil
_, err := f.WriteTo(bufio.NewWriter(&buf))
assert.EqualError(t, err, "zip: FileHeader.Name too long")
}
// Test StreamsWriter err
{
f := File{}
buf := bytes.Buffer{}
f.XLSX = make(map[string][]byte)
f.XLSX["s"] = nil
f.streams = make(map[string]*StreamWriter)
file, _ := os.Open("123")
f.streams["s"] = &StreamWriter{rawData: bufferedWriter{tmp: file}}
_, err := f.WriteTo(bufio.NewWriter(&buf))
assert.Nil(t, err)
}
}