feat: improved virtual text for assignments

Author: YaroSpace

lua/lua-console/utils.lua

@@ -125,9 +125,9 @@ 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)
+----Determines if there is an assigment on the line and returns its value
+----@param line string[]
+local get_line_assignment = function(line)
   if not line or #line == 0 then return end
 
   local lhs = line[1]:match('^(.-)%s*=')
@@ -140,39 +140,29 @@ local get_assignment = function(line)
   end
 end
 
-local print_buffer = {}
-
----@param buf number
----@param lines string[] Text to append to current buffer after current selection
-local append_current_buffer = function(buf, lines)
-  if not lines or #lines == 0 then return end
-
+local get_last_assignment = function()
+  local ctx = get_ctx()
   local lnum = vim.fn.line('.')
-  local prefix = config.buffer.result_prefix
-  local empty_results = { 'nil', '', '""', "''" }
-
-  local virtual_text
-  local line = lines[#lines]
 
-  if vim.tbl_contains(empty_results, line) then
-    table.remove(lines)
-    virtual_text = line
-  end
+  local last_var = ctx._last_assignment
+  local last_val = ctx[ctx._last_assignment]
 
-  local assignment_value = get_assignment(vim.fn.getbufline(buf, lnum, lnum))
-  if assignment_value ~= nil then virtual_text = assignment_value end
+  if not last_var then return end
 
-  if virtual_text then show_virtual_text(buf, 3, prefix .. virtual_text, lnum - 1, 'eol', 'Comment') end
+  local line
+  local offset = 0
 
-  if #lines == 0 then return end
+  for i = lnum - 1, 0, -1 do
+    line = vim.api.nvim_buf_get_lines(0, i, i + 1, false)[1]
 
-  lines[1] = prefix .. lines[1]
-  table.insert(lines, 1, '') -- insert an empty line
+    if line:match('^%s*' .. last_var .. '%s*=') then break end
+    offset = offset + 1
+  end
 
-  vim.api.nvim_buf_set_lines(buf, lnum, lnum, false, lines)
+  return last_var, last_val, offset
 end
 
----Pretty prints objects and appends to print_buffer
+---Pretty prints objects
 ---@param ... any[]
 ---@return string[]
 local pretty_print = function(...)
@@ -189,10 +179,47 @@ local pretty_print = function(...)
     result = result .. var_no .. vim.inspect(o)
   end
 
-  result = to_table(result)
-  vim.list_extend(print_buffer, result)
+  return to_table(result)
+end
+
+local print_buffer = {}
+
+local print_to_buffer = function(...)
+  local ret = pretty_print(...)
+  vim.list_extend(print_buffer, ret)
+end
+
+---@param buf number
+---@param lines string[] Text to append to current buffer after current selection
+local append_current_buffer = function(buf, lines)
+  if not lines or #lines == 0 then return end
+
+  local lnum = vim.fn.line('.')
+  local prefix = config.buffer.result_prefix
+  local empty_results = { 'nil', '', '""', "''" }
+
+  local virtual_text
+  local line = lines[#lines]
+
+  local last_assigned_var, last_assigned_val, last_assignment_offset = get_last_assignment()
+  if last_assigned_var then
+    virtual_text = to_string(pretty_print(last_assigned_val), '', true)
+    show_virtual_text(buf, 3, prefix .. virtual_text, lnum - last_assignment_offset - 1, 'eol', 'Comment')
+  end
+
+  if vim.tbl_contains(empty_results, line) then
+    table.remove(lines)
+
+    virtual_text = get_line_assignment(vim.fn.getbufline(buf, lnum, lnum)) or line
+    show_virtual_text(buf, 4, prefix .. virtual_text, lnum - 1, 'eol', 'Comment')
+  end
+
+  if #lines == 0 then return end
 
-  -- return result
+  lines[1] = prefix .. lines[1]
+  table.insert(lines, 1, '') -- insert an empty line
+
+  vim.api.nvim_buf_set_lines(buf, lnum, lnum, false, lines)
 end
 
 local function remove_empty_lines(tbl)
@@ -247,19 +274,26 @@ function get_ctx(buf)
   local ctx = lc.ctx[buf]
   if config.buffer.preserve_context and ctx then return ctx end
 
-  local env, mt = {}, {}
+  local env, mt, values = {}, {}, {}
 
   mt = {
-    print = pretty_print,
+    print = print_to_buffer,
     _ctx = function()
-      return vim.tbl_extend('force', {}, env)
+      return vim.deepcopy(values)
     end,
     _ctx_clear = function()
       lc.ctx[buf] = nil
     end,
     __index = function(_, key)
-      return mt[key] and mt[key] or _G[key]
+      return mt[key] and mt[key] or values[key] or _G[key]
+    end,
+    __newindex = function(_, k, v)
+      values[k] = v
+      mt._last_assignment = k
     end,
+    _reset_last_assignment = function()
+      mt._last_assignment = nil
+    end
   }
 
   lc.ctx[buf] = env

spec/unit/utils_spec.lua

@@ -479,6 +479,26 @@ describe('lua-console.utils', function()
         assert.has_string(result, expected)
       end)
 
+      it('shows value of the last assignment on the corresponging 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.result_prefix .. '5'
+
+        result = h.get_virtual_text(buf, 0, 0)
+        assert.has_string(result, expected)
+      end)
+
       it('gives syntax error message', function()
         vim.api.nvim_win_set_cursor(win, { 1, 0 })