fix: save_extra_cmds only works with single string

`Lib.run_hook_cmds` returns a table of tables (one for reach hook cmd
that was run). In `save_extra_cmds`, `vim.fn.writefile` expects a table
strings and if any of those strings contain newlines, they're
translated into nul characters. That means that there was no way to
return more than one line from a `save_extra_cmds` hook function.

To fix, we combine the results across all `save_extra_cmds` hooks and we
now support a `save_extra_cmds` hook returning either a string (possibly
with newlines) or table. If a hook returns a string then any newlines
will be used to split the string into individual line entries in the
table passed to `vim.fn.writefile`. If a hook returns a table, then each
of the values of the table will be added to the table passed to
`vim.fn.writefile` (the same newline splitting will also happen)

README.md

@@ -281,7 +281,7 @@ require('auto-session').setup({
 Command hooks exist in the format: {hook_name}
 
 - `{pre_save}`: executes _before_ a session is saved
-- `{save_extra}`: executes _after_ a session is saved, return string will save to `*x.vim`, reference `:help mks`
+- `{save_extra}`: executes _after_ a session is saved, saves returned string or table to `*x.vim`, reference `:help mks`
 - `{post_save}`: executes _after_ a session is saved
 - `{pre_restore}`: executes _before_ a session is restored
 - `{post_restore}`: executes _after_ a session is restored

lua/auto-session/init.lua

@@ -355,12 +355,19 @@ end
 ---@return boolean Returns whether extra commands were saved
 local function save_extra_cmds_new(session_path)
   local data = AutoSession.run_cmds "save_extra"
-  if not data then
+  local extra_file = string.gsub(session_path, "%.vim$", "x.vim")
+
+  -- data is a table of strings or tables, one for each hook function
+  -- need to combine them all here into a single table of strings
+  local data_to_write = Lib.flatten_table_and_split_strings(data)
+
+  if not data_to_write or vim.tbl_isempty(data_to_write) then
+    -- Have to delete the file just in case there's an old file from a previous save
+    vim.fn.delete(extra_file)
     return false
   end
 
-  local extra_file = string.gsub(session_path, "%.vim$", "x.vim")
-  if vim.fn.writefile(data, extra_file) ~= 0 then
+  if vim.fn.writefile(data_to_write, extra_file) ~= 0 then
     return false
   end
 

lua/auto-session/lib.lua

@@ -563,4 +563,45 @@ function Lib.run_hook_cmds(cmds, hook_name)
   return results
 end
 
+---Split any strings on newlines and add each one to the output table
+---Also flatten any embedded tables and and their values into the root table
+---(non recursive so only one level deep)
+---@param input table|nil
+---@return table The flattened table
+function Lib.flatten_table_and_split_strings(input)
+  local output = {}
+
+  if not input then
+    return output
+  end
+
+  local function add_value_to_output(value)
+    Lib.logger.debug("value: ", value)
+    if value == nil then
+      return
+    end
+
+    local value_type = type(value)
+    if value_type == "number" then
+      table.insert(output, value)
+    elseif value_type == "string" then
+      for s in value:gmatch "[^\r\n]+" do
+        table.insert(output, s)
+      end
+    end
+  end
+
+  for _, value in pairs(input) do
+    if type(value) == "table" then
+      for _, subvalue in pairs(value) do
+        add_value_to_output(subvalue)
+      end
+    else
+      add_value_to_output(value)
+    end
+  end
+
+  return output
+end
+
 return Lib

tests/extra_sesssion_commands_spec.lua

@@ -6,11 +6,15 @@ describe("Config with extra session commands", function()
   local save_extra_cmds_called = false
   local as = require "auto-session"
   local Lib = require "auto-session.lib"
+  -- WARN: this test calls setup again later to change save_extra_cmds
   as.setup {
     save_extra_cmds = {
       function()
         save_extra_cmds_called = true
-        return [[echo "hello world"]]
+        return [[
+        lua vim.g.extraCmdsTest = 1
+        lua vim.g.extraCmdsTest2 = 2
+        ]]
       end,
     },
     -- log_level = "debug",
@@ -36,7 +40,78 @@ describe("Config with extra session commands", function()
     TL.assertSessionHasFile(TL.default_session_path, TL.test_file)
 
     -- Make sure extra commands are there
-    assert.True(TL.fileHasString(default_extra_cmds_path, 'echo \\"hello world\\"'))
+    assert.True(TL.fileHasString(default_extra_cmds_path, "lua vim.g.extraCmdsTest = 1"))
+  end)
+
+  it("can restore a default session with extra commands", function()
+    vim.g.extraCmdsTest = 0
+    vim.g.extraCmdsTest2 = 0
+
+    assert.True(as.RestoreSession())
+
+    assert.True(vim.g.extraCmdsTest == 1)
+    assert.True(vim.g.extraCmdsTest2 == 2)
+  end)
+
+  it("can clear x.vim if there are no extra commands", function()
+    -- make sure the file is there now
+    assert.equals(1, vim.fn.filereadable(default_extra_cmds_path))
+
+    -- remove the handler
+    as.setup {
+      save_extra_cmds = nil,
+    }
+
+    -- generate default session
+    assert.True(as.AutoSaveSession())
+
+    -- Make sure the session was created
+    assert.equals(1, vim.fn.filereadable(TL.default_session_path))
+
+    -- make sure the extra commands file was removed
+    assert.equals(0, vim.fn.filereadable(default_extra_cmds_path))
+  end)
+
+  TL.clearSessionFilesAndBuffers()
+
+  it("can save a default session with extra commands in a table", function()
+    vim.cmd("e " .. TL.test_file)
+
+    save_extra_cmds_called = false
+
+    as.setup {
+      save_extra_cmds = {
+        function()
+          save_extra_cmds_called = true
+          return { "lua vim.g.extraCmdsTest = 1", "lua vim.g.extraCmdsTest2 = 2" }
+        end,
+      },
+    }
+
+    -- generate default session
+    assert.True(as.AutoSaveSession())
+
+    -- Make sure the session was created
+    assert.equals(1, vim.fn.filereadable(TL.default_session_path))
+    assert.equals(1, vim.fn.filereadable(default_extra_cmds_path))
+
+    assert.True(save_extra_cmds_called)
+
+    -- Make sure the session has our buffer
+    TL.assertSessionHasFile(TL.default_session_path, TL.test_file)
+
+    -- Make sure extra commands are there
+    assert.True(TL.fileHasString(default_extra_cmds_path, "lua vim.g.extraCmdsTest = 1"))
+  end)
+
+  it("can restore a default session with extra commands", function()
+    vim.g.extraCmdsTest = 0
+    vim.g.extraCmdsTest2 = 0
+
+    assert.True(as.RestoreSession())
+
+    assert.True(vim.g.extraCmdsTest == 1)
+    assert.True(vim.g.extraCmdsTest2 == 2)
   end)
 
   local session_name = "x"
@@ -61,7 +136,7 @@ describe("Config with extra session commands", function()
     TL.assertSessionHasFile(session_path, TL.test_file)
 
     -- Make sure extra commands are there
-    assert.True(TL.fileHasString(extra_cmds_path, 'echo \\"hello world\\"'))
+    assert.True(TL.fileHasString(extra_cmds_path, "lua vim.g.extraCmdsTest = 1"))
   end)
 
   it("can correctly differentiate x.vim session and xx.vim custom commands", function()

tests/lib_spec.lua

@@ -183,4 +183,36 @@ describe("Lib / Helper functions", function()
     assert.equals("", Lib.current_session_name())
     assert.equals("", Lib.current_session_name(true))
   end)
+
+  it("flatten_table_and_split_string() works with strings and tables", function()
+    local data = { "one\ntwo\nthree\n" }
+    local output = Lib.flatten_table_and_split_strings(data)
+
+    assert.equals(#output, 3)
+    assert.equals(output[1], "one")
+    assert.equals(output[2], "two")
+    assert.equals(output[3], "three")
+
+    data = { { "a", "b", "c\nd" }, "e", "f\ng" }
+    output = Lib.flatten_table_and_split_strings(data)
+
+    assert.equals(#output, 7)
+    assert.equals(output[1], "a")
+    assert.equals(output[2], "b")
+    assert.equals(output[3], "c")
+    assert.equals(output[4], "d")
+    assert.equals(output[5], "e")
+    assert.equals(output[6], "f")
+    assert.equals(output[7], "g")
+
+    data = {}
+    output = Lib.flatten_table_and_split_strings(data)
+
+    assert.True(vim.tbl_isempty(output))
+
+    data = { {} }
+    output = Lib.flatten_table_and_split_strings(data)
+
+    assert.True(vim.tbl_isempty(output))
+  end)
 end)