Skip to content

Commit

Permalink
qax-os#1402 Get CountRows from sheet
Browse files Browse the repository at this point in the history
  • Loading branch information
ivolkoff committed Nov 25, 2022
1 parent dde6b9c commit ea542a0
Show file tree
Hide file tree
Showing 2 changed files with 135 additions and 4 deletions.
87 changes: 83 additions & 4 deletions rows.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"io"
"math"
"os"
"regexp"
"strconv"

"github.com/mohae/deepcopy"
Expand Down Expand Up @@ -211,6 +212,15 @@ func (err ErrSheetNotExist) Error() string {
return fmt.Sprintf("sheet %s does not exist", err.SheetName)
}

// ErrCountRows defines an error of count rows
type ErrCountRows struct {
err error
}

func (err ErrCountRows) Error() string {
return fmt.Sprintf("wrong count rows: %s", err.err.Error())
}

// rowXMLIterator defined runtime use field for the worksheet row SAX parser.
type rowXMLIterator struct {
err error
Expand All @@ -237,6 +247,56 @@ func (rows *Rows) rowXMLHandler(rowIterator *rowXMLIterator, xmlElement *xml.Sta
}
}

// CountRows returns the number of rows in the worksheet.
// if return -1, that row not found
func (f *File) CountRows(sheet string) (int64, error) {
name, ok := f.getSheetXMLPath(sheet)
if !ok {
return -1, ErrSheetNotExist{sheet}
}

needClose, reader, tempFile, size, err := f.contentReader(name)
if err != nil {
return -1, ErrCountRows{fmt.Errorf("content reader: %v", err)}
}
if needClose && err == nil {
defer tempFile.Close()
}

var contentSize int64 = 1024
var content = make([]byte, contentSize)
var start int64
if size-contentSize < 0 {
start = 0
} else {
start = size - contentSize
}

if _, err = reader.ReadAt(content, start); err != nil && err != io.EOF {
return -1, ErrCountRows{fmt.Errorf("read at: %v", err)}
}

indexStart := bytes.LastIndex(content, []byte(`<row`))
if indexStart == -1 {
return 0, ErrCountRows{fmt.Errorf("not found row tag")}
}

indexStop := bytes.Index(content[indexStart:], []byte(`>`))
if indexStop == -1 {
return -1, ErrCountRows{fmt.Errorf("not found end row")}
}
indexStop = indexStart + indexStop

rFind := regexp.MustCompile(`r="(\d+)"`).Find(content[indexStart:indexStop])
if len(rFind) == 0 {
return -1, ErrCountRows{fmt.Errorf("not found row number")}
}

countStr := string(rFind[3 : len(rFind)-1])

return strconv.ParseInt(countStr, 10, 64)
}

// Rows returns a rows iterator, used for streaming reading data for a
// worksheet with a large data. This function is concurrency safe. For
// example:
Expand Down Expand Up @@ -326,19 +386,38 @@ func (f *File) getFromStringItem(index int) string {
return f.getFromStringItem(index)
}

// xmlDecoder creates XML decoder by given path in the zip from memory data
type ReaderContent interface {
io.Reader
io.ReaderAt
}

// contentReader returns reader by given path in the zip from memory data
// or system temporary file.
func (f *File) xmlDecoder(name string) (bool, *xml.Decoder, *os.File, error) {
func (f *File) contentReader(name string) (bool, ReaderContent, *os.File, int64, error) {
var (
content []byte
err error
tempFile *os.File
)
if content = f.readXML(name); len(content) > 0 {
return false, f.xmlNewDecoder(bytes.NewReader(content)), tempFile, err
return false, bytes.NewReader(content), tempFile, int64(len(content)), err
}

tempFile, err = f.readTemp(name)
return true, f.xmlNewDecoder(tempFile), tempFile, err

fileStat, err := tempFile.Stat()
if err != nil {
return true, tempFile, tempFile, 0, fmt.Errorf("failed to get file stat: %w", err)
}

return true, tempFile, tempFile, fileStat.Size(), err
}

// xmlDecoder creates XML decoder by given path in the zip from memory data
// or system temporary file.
func (f *File) xmlDecoder(name string) (bool, *xml.Decoder, *os.File, error) {
needClose, reader, tempFile, _, err := f.contentReader(name)
return needClose, f.xmlNewDecoder(reader), tempFile, err
}

// SetRowHeight provides a function to set the height of a single row. For
Expand Down
52 changes: 52 additions & 0 deletions rows_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1113,3 +1113,55 @@ func trimSliceSpace(s []string) []string {
}
return s
}

func TestFile_CountRows(t *testing.T) {
type fields struct {
filename string
}
tests := []struct {
name string
fields fields
want int64
wantErr assert.ErrorAssertionFunc
}{{
name: "BadWorkbook.xlsx",
fields: fields{filename: filepath.Join("test", "BadWorkbook.xlsx")},
want: -1,
wantErr: func(t assert.TestingT, err error, _ ...interface{}) bool {
return assert.Error(t, err)
},
}, {
name: "Book1.xlsx",
fields: fields{filename: filepath.Join("test", "Book1.xlsx")},
want: 22,
wantErr: func(t assert.TestingT, err error, _ ...interface{}) bool {
return assert.NoError(t, err)
},
}, {
name: "CalcChain.xlsx",
fields: fields{filename: filepath.Join("test", "CalcChain.xlsx")},
want: 1,
wantErr: func(t assert.TestingT, err error, _ ...interface{}) bool {
return assert.NoError(t, err)
},
}, {
name: "SharedStrings.xlsx",
fields: fields{filename: filepath.Join("test", "SharedStrings.xlsx")},
want: 1,
wantErr: func(t assert.TestingT, err error, _ ...interface{}) bool {
return assert.NoError(t, err)
},
}}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
f, err := OpenFile(tt.fields.filename)
assert.NoError(t, err)
firstSheet := f.GetSheetName(0)
got, err := f.CountRows(firstSheet)
if !tt.wantErr(t, err, "CountRows") {
return
}
assert.Equal(t, tt.want, got, "CountRows")
})
}
}

0 comments on commit ea542a0

Please sign in to comment.