feat: show result of assignments as virtual text

lua/lua-console/utils.lua

@@ -1,7 +1,21 @@
 local config = require('lua-console.config')
+local get_ctx, eval_lua
 
-local to_string = function(tbl)
-  return table.concat(tbl or {}, '\n')
+local to_string = function(tbl, sep, trim)
+  tbl = tbl or {}
+  sep = sep or '\n'
+
+  local line = table.concat(tbl, sep)
+  local patterns = { '\\r', '\\t', '\\n' }
+
+  if trim then
+    for _, pat in ipairs(patterns) do
+      line = line:gsub(pat, '')
+    end
+    line = line:gsub("([\"'])%s+", "%1"):gsub("%s+([\"'])", "%1"):gsub("%s%s+", ' ')
+  end
+
+  return line
 end
 
 local to_table = function(str)
@@ -12,13 +26,20 @@ local pack = function(...)
   return { ... }
 end
 
-local show_virtual_text = function(buf, id, text, line, position, highlight)
+---Shows virtual text in the buffer
+---@param buf number buffer
+---@param id number namespace id
+---@param text string text to show
+---@param lnum number line number
+---@param position string virtual text position
+---@param highlight string higlight group
+local show_virtual_text = function(buf, id, text, lnum, position, highlight)
   local ns = vim.api.nvim_create_namespace('Lua-console')
   local ext_mark = vim.api.nvim_buf_get_extmark_by_id(0, ns, id, {})
 
   if #ext_mark > 0 then vim.api.nvim_buf_del_extmark(0, ns, id) end
 
-  vim.api.nvim_buf_set_extmark(buf, ns, line, 0, {
+  vim.api.nvim_buf_set_extmark(buf, ns, lnum, 0, {
     id = id,
     virt_text = { { text, highlight } },
     virt_text_pos = position,
@@ -92,21 +113,45 @@ local get_path_lnum = function(path)
   return path, lnum
 end
 
+---Determines if there is an assigment on the line and returns its value
+---@param line string[]
+local get_assignment = function(line)
+  local lhs = line[1]:match('^(.-)%s*=')
+  local ret
+
+  if lhs then
+    ret = eval_lua({ lhs })
+    local is_error = ret[1]:find('[string "Lua console: "]', 1, true) -- means we could not evaluate the lhs
+    return not is_error and to_string(ret, '', true) or nil
+  end
+end
+
 local print_buffer = {}
 
 ---@param lines string[] Text to append to buffer after current selection
 local append_current_buffer = function(lines)
   local buf = vim.fn.bufnr()
   local lnum = math.max(vim.fn.line('.'), vim.fn.line('v'))
+  local prepend = config.buffer.prepend_result_with
+
+  local virtual_text
+  if lines[#lines] == 'nil' then
+    table.remove(lines)
+    virtual_text = 'nil'
+  end
 
-  lines[1] = config.buffer.prepend_result_with .. lines[1]
+  local assignment_value = get_assignment(vim.fn.getbufline(buf, lnum, lnum))
+  if assignment_value ~= nil then virtual_text = assignment_value end
 
-  if #lines == 1 and lines[1]:find('nil') then
-    show_virtual_text(buf, 3, lines[1], lnum - 1, 'eol', 'Comment')
-    return
+  if virtual_text then
+    show_virtual_text(buf, 3, prepend .. virtual_text, lnum - 1, 'eol', 'Comment')
   end
 
+  if #lines == 0 then return end
+
+  lines[1] = prepend .. lines[1]
   table.insert(lines, 1, '') -- insert an empty line
+
   vim.api.nvim_buf_set_lines(buf, lnum, lnum, false, lines)
 end
 
@@ -171,7 +216,7 @@ local function add_return(tbl)
   return ret
 end
 
-local get_ctx = function(buf)
+get_ctx = function(buf)
   buf = buf or vim.fn.bufnr()
   Lua_console.ctx = Lua_console.ctx or {}
 
@@ -199,9 +244,9 @@ end
 
 --- Evaluates Lua code and returns pretty printed result with errors if any
 --- @param lines string[] table with lines of Lua code
---- @param ctx table environment to execute code in
+--- @param ctx? table environment to execute code in
 --- @return string[]
-local eval_lua = function(lines, ctx)
+eval_lua = function(lines, ctx)
   vim.validate({ lines = { lines, 'table'} })
 
   local lines_with_return = add_return(lines)
@@ -218,14 +263,15 @@ local eval_lua = function(lines, ctx)
 
   ---@cast code function
   local result = pack(xpcall(code, debug.traceback))
-  vim.api.nvim_buf_set_var(vim.fn.bufnr(), 'ctx', env)
+  local status, err = result[1], result[2]
 
-  if result[1] then
+  if status then
     table.remove(result, 1)
+
     if #result > 0 then pretty_print(unpack(result))
     else pretty_print(nil) end
   else
-    vim.list_extend(print_buffer, clean_stacktrace(result[2]))
+    vim.list_extend(print_buffer, clean_stacktrace(err))
   end
 
   return print_buffer

spec/unit/mappings_spec.lua

@@ -86,7 +86,6 @@ describe("lua-console.nvim - mappings", function()
 				=> [1] 1, [2] "A-10"
 				[1] 2, [2] "A-20"
 				[1] 3, [2] "A-30"
-				nil
 			]]
 
 			assert.has_string(result, expected)

spec/unit/utils_spec.lua

@@ -31,7 +31,7 @@ describe('lua-console.utils', function()
   		-- code
   	end)
 
-    pending('preserves context for different buffers', function()
+    pending('preserves context', function()
       -- code
     end)
 
@@ -253,7 +253,7 @@ describe('lua-console.utils', function()
   				a = 1,
   				b = 2
 				}
-				nil]]
+				]]
 
       result = eval_lua(code)
       assert.has_string(result, expected)
@@ -331,7 +331,7 @@ describe('lua-console.utils', function()
         h.set_buffer(buf, content)
       end)
 
-      it('evaluates lua in current buffer - single line', function()
+      it('single line', function()
         vim.api.nvim_win_set_cursor(win, { 1, 0 })
         utils.eval_code_in_buffer()
 
@@ -348,7 +348,7 @@ describe('lua-console.utils', function()
         assert.has_string(result, expected)
       end)
 
-      it('evaluates lua in current buffer - multiline', function()
+      it('multiline', function()
         vim.api.nvim_win_set_cursor(win, { 2, 0 })
         vim.cmd.exe("'normal V3j'")
 
@@ -371,15 +371,14 @@ describe('lua-console.utils', function()
 					[1] 8, [2] 40
 					[1] 9, [2] 45
 					[1] 10, [2] 50
-					nil
 					Some text here
   			]])
 
         result = h.get_buffer(buf)
         assert.has_string(result, expected)
       end)
 
-      it('evaluates lua in current buffer - shows nil as virtual text', function()
+      it('shows nil as virtual text', function()
         vim.api.nvim_win_set_cursor(win, { 1, 0 })
         h.send_keys('V2j')
 
@@ -397,6 +396,26 @@ describe('lua-console.utils', function()
         assert.has_string(result, expected)
       end)
 
+      it('shows value of the assignment on the last line as virtual text', function()
+        vim.api.nvim_win_set_cursor(win, { 1, 0 })
+        h.send_keys('V4j')
+
+        content = h.to_table([[
+					a = 5
+					for i = 1, 5 do
+						i = i + 5
+					end
+					vim.bo.filetype = tostring(a + 5) .. 'test'
+  			]])
+        h.set_buffer(buf, content)
+        utils.eval_code_in_buffer()
+
+        expected = config.buffer.prepend_result_with .. '"10test"'
+
+        result = h.get_virtual_text(buf, 4, 4)
+        assert.has_string(result, expected)
+      end)
+
       it('gives syntax error message', function()
         vim.api.nvim_win_set_cursor(win, { 1, 0 })