Merge remote-tracking branch 'myk002/quickfort_xlsx' into quickfort

lethosor · lethosor · commit d30296b9e97b · 2020-08-17T22:58:17.000-04:00
diff --git a/internal/quickfort/command.lua b/internal/quickfort/command.lua
@@ -27,8 +27,8 @@ local valid_command_args = utils.invert({
     '-quiet',
     'v',
     '-verbose',
-    's',
-    '-sheet',
+    'n',
+    '-name',
 })
 
 function do_command(in_args)
@@ -42,14 +42,15 @@ function do_command(in_args)
         qerror("expected <list_num> or <blueprint_name> parameter")
     end
     local list_num = tonumber(blueprint_name)
+    local sheet_name = nil
     if list_num then
-        blueprint_name = quickfort_list.get_blueprint_by_number(list_num)
+        blueprint_name, sheet_name =
+                quickfort_list.get_blueprint_by_number(list_num)
     end
-
     local args = utils.processArgs(in_args, valid_command_args)
     local quiet = args['q'] ~= nil or args['-quiet'] ~= nil
     local verbose = args['v'] ~= nil or args['-verbose'] ~= nil
-    local sheet = tonumber(args['s']) or tonumber(args['-sheet'])
+    sheet_name = sheet_name or args['n'] or args['-name']
 
     local cursor = guidm.getCursorPos()
     if command ~= 'orders' and not cursor then
@@ -60,7 +61,7 @@ function do_command(in_args)
     quickfort_common.verbose = verbose
 
     local filepath = quickfort_common.get_blueprint_filepath(blueprint_name)
-    local data = quickfort_parse.process_file(filepath, cursor)
+    local data = quickfort_parse.process_file(filepath, sheet_name, cursor)
     for zlevel, section_data_list in pairs(data) do
         for _, section_data in ipairs(section_data_list) do
             local modeline = section_data.modeline
diff --git a/internal/quickfort/list.lua b/internal/quickfort/list.lua
@@ -6,6 +6,7 @@ if not dfhack_flags.module then
 end
 
 local utils = require('utils')
+local xlsxreader = require('plugins.xlsxreader')
 local quickfort_common = reqscript('internal/quickfort/common')
 local quickfort_parse = reqscript('internal/quickfort/parse')
 
@@ -20,64 +21,109 @@ end
 
 local blueprint_cache = {}
 
-local function scan_blueprint(path)
+local function scan_csv_blueprint(path)
     local filepath = quickfort_common.get_blueprint_filepath(path)
     local mtime = dfhack.filesystem.mtime(filepath)
     if not blueprint_cache[path] or blueprint_cache[path].mtime ~= mtime then
         blueprint_cache[path] = {modeline=get_modeline(filepath), mtime=mtime}
     end
+    if not blueprint_cache[path].modeline then
+        print(string.format('skipping "%s": no #mode marker detected', path))
+    end
     return blueprint_cache[path].modeline
 end
 
-local blueprint_files = {}
+local function get_xlsx_sheet_modeline(xlsx_file, sheet_name)
+    local xlsx_sheet = xlsxreader.open_sheet(xlsx_file, sheet_name)
+    return dfhack.with_finalize(
+        function() xlsxreader.close_sheet(xlsx_sheet) end,
+        function()
+            local row_cells = xlsxreader.get_row(xlsx_sheet)
+            if not row_cells or #row_cells == 0 then return nil end
+            return quickfort_parse.parse_modeline(row_cells[1])
+        end
+    )
+end
+
+local function get_xlsx_file_sheet_infos(filepath)
+    local sheet_infos = {}
+    local xlsx_file = xlsxreader.open_xlsx_file(filepath)
+    if not xlsx_file then return sheet_infos end
+    return dfhack.with_finalize(
+        function() xlsxreader.close_xlsx_file(xlsx_file) end,
+        function()
+            for _, sheet_name in ipairs(xlsxreader.list_sheets(xlsx_file)) do
+                local modeline = get_xlsx_sheet_modeline(xlsx_file, sheet_name)
+                if modeline then
+                    table.insert(sheet_infos,
+                                 {name=sheet_name, modeline=modeline})
+                end
+            end
+            return sheet_infos
+        end
+    )
+end
+
+local function scan_xlsx_blueprint(path)
+    local filepath = quickfort_common.get_blueprint_filepath(path)
+    local mtime = dfhack.filesystem.mtime(filepath)
+    if blueprint_cache[path] and blueprint_cache[path].mtime == mtime then
+        return blueprint_cache[path].sheet_infos
+    end
+    local sheet_infos = get_xlsx_file_sheet_infos(filepath)
+    if #sheet_infos == 0 then
+        print(string.format(
+                'skipping "%s": no sheet with #mode markers detected', path))
+    end
+    blueprint_cache[path] = {sheet_infos=sheet_infos, mtime=mtime}
+    return sheet_infos
+end
+
+local blueprints = {}
 
 local function scan_blueprints()
     local paths = dfhack.filesystem.listdir_recursive(
         quickfort_common.settings['blueprints_dir'].value, nil, false)
-    blueprint_files = {}
-    local library_files = {}
+    blueprints = {}
+    local library_blueprints = {}
     for _, v in ipairs(paths) do
-        if not v.isdir and
-                (string.find(v.path, '[.]csv$') or
-                 string.find(v.path, '[.]xlsx$')) then
-            if string.find(v.path, '[.]xlsx$') then
-                print(string.format(
-                        'skipping "%s": .xlsx files not supported yet', v.path))
-                goto skip
+        local is_library = string.find(v.path, '^library/') ~= nil
+        local target_list = blueprints
+        if is_library then target_list = library_blueprints end
+        if not v.isdir and string.find(v.path:lower(), '[.]csv$') then
+            local modeline = scan_csv_blueprint(v.path)
+            if modeline then
+                table.insert(target_list,
+                        {path=v.path, modeline=modeline, is_library=is_library})
             end
-            local modeline = scan_blueprint(v.path)
-            if not modeline then
-                print(string.format(
-                        'skipping "%s": no #mode marker detected', v.path))
-                goto skip
+        elseif not v.isdir and string.find(v.path:lower(), '[.]xlsx$') then
+            local sheet_infos = scan_xlsx_blueprint(v.path)
+            if #sheet_infos > 0 then
+                for _, sheet_info in ipairs(sheet_infos) do
+                    table.insert(target_list,
+                            {path=v.path,
+                             sheet_name=sheet_info.name,
+                             modeline=sheet_info.modeline,
+                             is_library=is_library})
+                end
             end
-            if string.find(v.path, '^library/') ~= nil then
-                table.insert(
-                    library_files,
-                    {path=v.path, modeline=modeline, is_library=true})
-            else
-                table.insert(
-                    blueprint_files,
-                    {path=v.path, modeline=modeline, is_library=false})
-            end
-            ::skip::
         end
     end
     -- tack library files on to the end so user files are contiguous
-    for i=1, #library_files do
-        blueprint_files[#blueprint_files + 1] = library_files[i]
+    for i=1, #library_blueprints do
+        blueprints[#blueprints + 1] = library_blueprints[i]
     end
 end
 
 function get_blueprint_by_number(list_num)
-    if #blueprint_files == 0 then
+    if #blueprints == 0 then
         scan_blueprints()
     end
-    local blueprint_file = blueprint_files[list_num]
-    if not blueprint_file then
+    local blueprint = blueprints[list_num]
+    if not blueprint then
         qerror(string.format('invalid list index: %d', list_num))
     end
-    return blueprint_file.path
+    return blueprint.path, blueprint.sheet_name
 end
 
 local valid_list_args = utils.invert({
@@ -89,8 +135,12 @@ function do_list(in_args)
     local args = utils.processArgs(in_args, valid_list_args)
     local show_library = args['l'] ~= nil or args['-library'] ~= nil
     scan_blueprints()
-    for i, v in ipairs(blueprint_files) do
+    for i, v in ipairs(blueprints) do
         if show_library or not v.is_library then
+            local sheet_spec = ''
+            if v.sheet_name then
+                sheet_spec = string.format(' -n "%s"', v.sheet_name)
+            end
             local comment = ')'
             if #v.modeline.comment > 0 then
                 comment = string.format(': %s)', v.modeline.comment)
@@ -100,8 +150,8 @@ function do_list(in_args)
                 start_comment = string.format('; cursor start: %s',
                                               v.modeline.start_comment)
             end
-            print(string.format('%d) "%s" (%s%s%s',
-                                i, v.path, v.modeline.mode, comment,
+            print(string.format('%d) "%s"%s (%s%s%s',
+                                i, v.path, sheet_spec, v.modeline.mode, comment,
                                 start_comment))
         end
     end
diff --git a/internal/quickfort/parse.lua b/internal/quickfort/parse.lua
@@ -5,6 +5,7 @@ if not dfhack_flags.module then
     qerror('this script cannot be called directly')
 end
 
+local xlsxreader = require('plugins.xlsxreader')
 local quickfort_common = reqscript('internal/quickfort/common')
 
 local function trim_and_insert(tokens, token)
@@ -13,6 +14,7 @@ local function trim_and_insert(tokens, token)
 end
 
 -- adapted from example on http://lua-users.org/wiki/LuaCsv
+-- returns a list of strings corresponding to the text in the cells in the row
 function tokenize_csv_line(line)
     line = string.gsub(line, '[\r\n]*$', '')
     local tokens = {}
@@ -108,16 +110,62 @@ local function make_cell_label(col_num, row_num)
     return get_col_name(col_num) .. tostring(math.floor(row_num))
 end
 
+local function read_csv_line(ctx)
+    local line = ctx.csv_file:read()
+    if not line then return nil end
+    return tokenize_csv_line(line)
+end
+
+local function cleanup_csv_ctx(ctx)
+    ctx.csv_file:close()
+end
+
+local function read_xlsx_line(ctx)
+    return xlsxreader.get_row(ctx.xlsx_sheet)
+end
+
+local function cleanup_xslx_ctx(ctx)
+    xlsxreader.close_sheet(ctx.xlsx_sheet)
+    xlsxreader.close_xlsx_file(ctx.xlsx_file)
+end
+
+local function init_reader_ctx(filepath, sheet_name)
+    local reader_ctx = {}
+    if string.find(filepath:lower(), '[.]csv$') then
+        local file = io.open(filepath)
+        if not file then
+            qerror(string.format('failed to open blueprint file: "%s"',
+                                 filepath))
+        end
+        reader_ctx.csv_file = file
+        reader_ctx.get_row_tokens = read_csv_line
+        reader_ctx.cleanup = cleanup_csv_ctx
+    else
+        reader_ctx.xlsx_file = xlsxreader.open_xlsx_file(filepath)
+        if not reader_ctx.xlsx_file then
+            qerror(string.format('failed to open blueprint file: "%s"',
+                                 filepath))
+        end
+        -- open_sheet succeeds even if the sheet cannot be found; we need to
+        -- check that when we try to read
+        reader_ctx.xlsx_sheet =
+                xlsxreader.open_sheet(reader_ctx.xlsx_file, sheet_name)
+        reader_ctx.get_row_tokens = read_xlsx_line
+        reader_ctx.cleanup = cleanup_xslx_ctx
+    end
+    return reader_ctx
+end
+
 -- returns a grid representation of the current section, the number of lines
 -- read from the input, and the next z-level modifier, if any. See process_file
 -- for grid format.
-local function process_section(file, start_line_num, start_coord)
+local function process_section(reader_ctx, start_line_num, start_coord)
     local grid = {}
     local y = start_coord.y
     while true do
-        local line = file:read()
-        if not line then return grid, y-start_coord.y end
-        for i, v in ipairs(tokenize_csv_line(line)) do
+        local row_tokens = reader_ctx.get_row_tokens(reader_ctx)
+        if not row_tokens then return grid, y-start_coord.y end
+        for i, v in ipairs(row_tokens) do
             if i == 1 then
                 if v == '#<' then return grid, y-start_coord.y, 1 end
                 if v == '#>' then return grid, y-start_coord.y, -1 end
@@ -135,32 +183,22 @@ local function process_section(file, start_line_num, start_coord)
     end
 end
 
---[[
-returns the following logical structure:
-  map of target map z coordinate ->
-    list of {modeline, grid} tables
-Where the structure of modeline is defined as per parse_modeline and grid is a:
-  map of target y coordinate ->
-    map of target map x coordinate ->
-      {cell=spreadsheet cell, text=text from spreadsheet cell}
-Map keys are numbers, and the keyspace is sparse -- only elements that have
-contents are non-nil.
-]]
-function process_file(filepath, start_cursor_coord)
-    local file = io.open(filepath)
-    if not file then
-        qerror(string.format('failed to open blueprint file: "%s"', filepath))
+function process_sections(reader_ctx, filepath, sheet_name, start_cursor_coord)
+    local row_tokens = reader_ctx.get_row_tokens(reader_ctx)
+    if not row_tokens then
+        qerror(string.format(
+                'sheet with name: "%s" in file "%s" empty or not found',
+                sheet_name, filepath))
     end
-    local line = file:read()
-    local modeline = parse_modeline(tokenize_csv_line(line)[1])
+    local modeline = parse_modeline(row_tokens[1])
     local cur_line_num = 2
     local x = start_cursor_coord.x - modeline.startx + 1
     local y = start_cursor_coord.y - modeline.starty + 1
     local z = start_cursor_coord.z
     local zlevels = {}
     while true do
         local grid, num_section_rows, zmod =
-                process_section(file, cur_line_num, xyz2pos(x, y, z))
+                process_section(reader_ctx, cur_line_num, xyz2pos(x, y, z))
         for _, _ in pairs(grid) do
             -- apparently, the only way to tell if a sparse array is not empty
             if not zlevels[z] then zlevels[z] = {} end
@@ -171,7 +209,27 @@ function process_file(filepath, start_cursor_coord)
         cur_line_num = cur_line_num + num_section_rows + 1
         z = z + zmod
     end
-    file:close()
     return zlevels
 end
 
+--[[
+returns the following logical structure:
+  map of target map z coordinate ->
+    list of {modeline, grid} tables
+Where the structure of modeline is defined as per parse_modeline and grid is a:
+  map of target y coordinate ->
+    map of target map x coordinate ->
+      {cell=spreadsheet cell, text=text from spreadsheet cell}
+Map keys are numbers, and the keyspace is sparse -- only elements that have
+contents are non-nil.
+]]
+function process_file(filepath, sheet_name, start_cursor_coord)
+    local reader_ctx = init_reader_ctx(filepath, sheet_name)
+    return dfhack.with_finalize(
+        function() reader_ctx.cleanup(reader_ctx) end,
+        function()
+            return process_sections(reader_ctx, filepath, sheet_name,
+                                    start_cursor_coord)
+        end
+    )
+end
diff --git a/quickfort.lua b/quickfort.lua
@@ -39,10 +39,10 @@ Usage:
     subfolder are not shown. Specify ``-l`` to include library blueprints.
 **quickfort <command> <list_num> [<options>]**
     Applies the blueprint with the number from the list command.
-**quickfort <command> <filename> [-s|--sheet <sheet_num>] [<options>]**
+**quickfort <command> <filename> [-n|--name <sheet_name>] [<options>]**
     Applies the blueprint from the named file. If it is an ``.xlsx`` file,
-    the ``-s`` (or ``--sheet``) parameter is required to identify the sheet
-    number. The first sheet is ``-s 1``.
+    the ``-n`` (or ``--name``) parameter can identify the sheet name. If the
+    sheet name is not specified, the first sheet is used.
 
 **<command>** can be one of:
 
