Lua class wrappers for the xlsxreader plugin API

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
     )

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

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<std::string>();
         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);
     }