diff --git a/docs/Lua API.rst b/docs/Lua API.rst index 764e56dd6..c48d522f6 100644 --- a/docs/Lua API.rst +++ b/docs/Lua API.rst @@ -4305,41 +4305,44 @@ xlsxreader ========== Utility functions to facilitate reading .xlsx spreadsheets. It provides the -following API methods: +following low-level API methods: - ``file_handle open_xlsx_file(filename)`` - ``close_xlsx_file(file_handle)`` - ``sheet_names list_sheets(file_handle)`` - ``sheet_handle open_sheet(file_handle, sheet_name)`` - ``close_sheet(sheet_handle)`` - - ``cell_strings get_row(sheet_handle)`` + - ``cell_strings get_row(sheet_handle, max_tokens)`` The ``max_tokens`` + parameter is optional. If set to a number > 0, it limits the number of cells + read and returned for the row. - Example:: + Lua users would be better off using the Lua class wrappers, though. For + example:: local xlsxreader = require('plugins.xlsxreader') - local function dump_sheet(xlsx_file, sheet_name) - print('reading sheet: '..sheet_name) - local xlsx_sheet = xlsxreader.open_sheet(xlsx_file, sheet_name) + local function dump_sheet(reader, sheet_name) + print('reading sheet: ' .. sheet_name) + local sheet_reader = reader:open_sheet(sheet_name) dfhack.with_finalize( - function () xlsxreader.close_sheet(xlsx_sheet) end, - function () - local row_cells = xlsxreader.get_row(xlsx_sheet) + function() sheet_reader:close() end, + function() + local row_cells = sheet_reader:get_row() while row_cells do printall(row_cells) - row_cells = xlsxreader.get_row(xlsx_sheet) + row_cells = sheet_reader:get_row() end end ) end - local filepath = "path/to/some_file.xlsx" - local xlsx_file = xlsxreader.open_xlsx_file(filepath) + local filepath = 'path/to/some_file.xlsx' + local reader = xlsxreader.XlsxioReader{filepath=filepath} dfhack.with_finalize( - function () xlsxreader.close_xlsx_file(xlsx_file) end, - function () - for _, sheet_name in ipairs(xlsxreader.list_sheets(xlsx_file)) do - dump_sheet(xlsx_file, sheet_name) + function() reader:close() end, + function() + for _,sheet_name in ipairs(reader:list_sheets()) do + dump_sheet(reader, sheet_name) end end ) diff --git a/plugins/lua/xlsxreader.lua b/plugins/lua/xlsxreader.lua index 323231915..b9827c39f 100644 --- a/plugins/lua/xlsxreader.lua +++ b/plugins/lua/xlsxreader.lua @@ -1,47 +1,55 @@ local _ENV = mkmodule('plugins.xlsxreader') ---[[ - - Native functions: - - * file_handle open_xlsx_file(filename) - * close_xlsx_file(file_handle) - * sheet_names list_sheets(file_handle) - - * sheet_handle open_sheet(file_handle, sheet_name) - * close_sheet(sheet_handle) - * cell_strings get_row(sheet_handle) - -Sample usage: - - local xlsxreader = require('plugins.xlsxreader') - - local function dump_sheet(xlsx_file, sheet_name) - print('reading sheet: '..sheet_name) - local xlsx_sheet = xlsxreader.open_sheet(xlsx_file, sheet_name) - dfhack.with_finalize( - function () xlsxreader.close_sheet(xlsx_sheet) end, - function () - local row_cells = xlsxreader.get_row(xlsx_sheet) - while row_cells do - printall(row_cells) - row_cells = xlsxreader.get_row(xlsx_sheet) - end - end - ) +XlsxioSheetReader = defclass(XlsxioSheetReader, nil) +XlsxioSheetReader.ATTRS{ + -- handle to the open sheet. required. + sheet_handle = DEFAULT_NIL, +} + +function XlsxioSheetReader:init() + if not self.sheet_handle then + error('XlsxSheetReader: sheet_handle is required') end - - local filepath = "path/to/some_file.xlsx" - local xlsx_file = xlsxreader.open_xlsx_file(filepath) - dfhack.with_finalize( - function () xlsxreader.close_xlsx_file(xlsx_file) end, - function () - for _, sheet_name in ipairs(xlsxreader.list_sheets(xlsx_file)) do - dump_sheet(xlsx_file, sheet_name) - end - end - ) - ---]] +end + +function XlsxioSheetReader:close() + close_sheet(self.sheet_handle) + self.sheet_handle = nil +end + +function XlsxioSheetReader:get_row(max_tokens) + return get_row(self.sheet_handle, max_tokens) +end + +XlsxioReader = defclass(XlsxioReader, nil) +XlsxioReader.ATTRS{ + -- full or relative path to the target .xlsx file. required. + filepath = DEFAULT_NIL, +} + +function XlsxioReader:init() + if not self.filepath then + error('XlsxReader: filepath is required') + end + self.xlsx_handle = open_xlsx_file(self.filepath) + if not self.xlsx_handle then + qerror(('failed to open "%s"'):format(self.filepath)) + end +end + +function XlsxioReader:close() + close_xlsx_file(self.xlsx_handle) + self.xlsx_handle = nil +end + +function XlsxioReader:list_sheets() + return list_sheets(self.xlsx_handle) +end + +-- if sheet_name is empty or nil, opens the first sheet +function XlsxioReader:open_sheet(sheet_name) + local sheet_handle = open_sheet(self.xlsx_handle, sheet_name) + return XlsxioSheetReader{sheet_handle=sheet_handle} +end return _ENV diff --git a/plugins/xlsxreader.cpp b/plugins/xlsxreader.cpp index 7f61d6ff5..d9a781f11 100644 --- a/plugins/xlsxreader.cpp +++ b/plugins/xlsxreader.cpp @@ -101,11 +101,13 @@ int list_sheets(lua_State *L) { return 1; } -// takes the sheet handle and returns a table-list of strings, or nil if we -// already processed the last row in the file. +// takes the sheet handle and returns a table-list of strings, limited to the +// requested number of tokens if specified and > 0, or nil if we already +// processed the last row in the file. int get_row(lua_State *L) { auto sheet_handle = (xlsx_sheet_handle *)get_xlsxreader_handle(L); CHECK_NULL_POINTER(sheet_handle->handle); + lua_Integer max_tokens = luaL_optinteger(L, 2, 0); bool ok = get_next_row(sheet_handle); if (!ok) { lua_pushnil(L); @@ -114,6 +116,8 @@ int get_row(lua_State *L) { auto cells = std::vector(); while (get_next_cell(sheet_handle, value)) { cells.push_back(value); + if (max_tokens > 0 && cells.size() >= max_tokens) + break; } Lua::PushVector(L, cells, true); }