feat: stream write to zip directly (#863)
This commit is contained in:
parent
2cfcf9eb5f
commit
24967a5c25
118
file.go
118
file.go
|
@ -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
|
||||
}
|
||||
|
|
43
file_test.go
43
file_test.go
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue