From a0a51c06c2fcddda925667142516c89777eb0c8e Mon Sep 17 00:00:00 2001
From: Folke Lemaitre <folke.lemaitre@gmail.com>
Date: Wed, 26 Jun 2024 21:38:28 +0200
Subject: [PATCH] feat: added `opts.headless` to control ansi output when
 running headless

---
 lua/lazy/core/config.lua      | 11 ++++++
 lua/lazy/manage/process.lua   |  6 +++
 lua/lazy/manage/task/init.lua | 59 ++++++++++++++++++++++++++++-
 lua/lazy/terminal.lua         | 71 +++++++++++++++++++++++++++++++++++
 tests/busted.lua              | 16 ++++----
 5 files changed, 153 insertions(+), 10 deletions(-)
 create mode 100644 lua/lazy/terminal.lua

diff --git a/lua/lazy/core/config.lua b/lua/lazy/core/config.lua
index a6d3c3a..89a89bd 100644
--- a/lua/lazy/core/config.lua
+++ b/lua/lazy/core/config.lua
@@ -130,6 +130,17 @@ M.defaults = {
       },
     },
   },
+  -- Output options for headless mode
+  headless = {
+    -- show the output from process commands like git
+    process = true,
+    -- show log messages
+    log = true,
+    -- show task start/end
+    task = true,
+    -- use ansi colors
+    colors = true,
+  },
   diff = {
     -- diff command <d> can be one of:
     -- * browser: opens the github compare view. Note that this is always mapped to <K> as well,
diff --git a/lua/lazy/manage/process.lua b/lua/lazy/manage/process.lua
index 5bb2716..7c568cb 100644
--- a/lua/lazy/manage/process.lua
+++ b/lua/lazy/manage/process.lua
@@ -48,6 +48,7 @@ local uv = vim.uv
 ---@field cwd? string
 ---@field on_line? fun(string)
 ---@field on_exit? fun(ok:boolean, output:string)
+---@field on_data? fun(string)
 ---@field timeout? number
 ---@field env? table<string,string>
 
@@ -145,6 +146,11 @@ function M.spawn(cmd, opts)
     assert(not err, err)
 
     if data then
+      if opts.on_data then
+        vim.schedule(function()
+          opts.on_data(data)
+        end)
+      end
       output = output .. data:gsub("\r\n", "\n")
       local lines = vim.split(vim.trim(output:gsub("\r$", "")):gsub("[^\n\r]+\r", ""), "\n")
 
diff --git a/lua/lazy/manage/task/init.lua b/lua/lazy/manage/task/init.lua
index 8c5f83f..d561151 100644
--- a/lua/lazy/manage/task/init.lua
+++ b/lua/lazy/manage/task/init.lua
@@ -1,5 +1,9 @@
 local Async = require("lazy.async")
+local Config = require("lazy.core.config")
 local Process = require("lazy.manage.process")
+local Terminal = require("lazy.terminal")
+
+local colors = Config.options.headless.colors
 
 ---@class LazyTaskDef
 ---@field skip? fun(plugin:LazyPlugin, opts?:TaskOptions):any?
@@ -96,6 +100,10 @@ function Task:_start(task)
   assert(not self:has_started(), "task already started")
   assert(not self:has_ended(), "task already done")
 
+  if Config.headless() and Config.options.headless.task then
+    self:log("Running task " .. self.name, vim.log.levels.INFO)
+  end
+
   self._started = vim.uv.hrtime()
   ---@async
   self._running = Async.run(function()
@@ -122,6 +130,27 @@ function Task:log(msg, level)
   ---@cast msg string
   table.insert(self._log, { msg = msg, level = level })
   vim.api.nvim_exec_autocmds("User", { pattern = "LazyRender", modeline = false })
+  if Config.headless() then
+    self:headless()
+  end
+end
+
+function Task:headless()
+  if not Config.options.headless.log then
+    return
+  end
+  local msg = self._log[#self._log]
+  if not msg or msg.level == vim.log.levels.TRACE then
+    return
+  end
+  local map = {
+    [vim.log.levels.ERROR] = Terminal.red,
+    [vim.log.levels.WARN] = Terminal.yellow,
+    [vim.log.levels.INFO] = Terminal.blue,
+  }
+  local color = Config.options.headless.colors and map[msg.level]
+  io.write(Terminal.prefix(color and color(msg.msg) or msg.msg, self:prefix()))
+  io.write("\n")
 end
 
 ---@param msg string|string[]
@@ -143,6 +172,10 @@ function Task:_done()
     return
   end
 
+  if Config.headless() and Config.options.headless.task then
+    local ms = math.floor(self:time() + 0.5)
+    self:log("Finished task " .. self.name .. " in " .. ms .. "ms", vim.log.levels.INFO)
+  end
   self._ended = vim.uv.hrtime()
   if self._opts.on_done then
     self._opts.on_done(self)
@@ -172,8 +205,12 @@ function Task:spawn(cmd, opts)
   local on_line = opts.on_line
   local on_exit = opts.on_exit
 
+  local headless = Config.headless() and Config.options.headless.process
+
   function opts.on_line(line)
-    self:log(line, vim.log.levels.TRACE)
+    if not headless then
+      return self:log(line, vim.log.levels.TRACE)
+    end
     if on_line then
       pcall(on_line, line)
     end
@@ -182,18 +219,36 @@ function Task:spawn(cmd, opts)
   local running = true
   ---@param output string
   function opts.on_exit(ok, output)
-    self:log(vim.trim(output), ok and vim.log.levels.DEBUG or vim.log.levels.ERROR)
+    if not headless then
+      self:log(vim.trim(output), ok and vim.log.levels.DEBUG or vim.log.levels.ERROR)
+    end
     if on_exit then
       pcall(on_exit, ok, output)
     end
     running = false
   end
+
+  if headless then
+    opts.on_data = function(data)
+      -- prefix with plugin name
+      local prefix = self:prefix()
+      io.write(Terminal.prefix(data, prefix))
+    end
+  end
   Process.spawn(cmd, opts)
   while running do
     coroutine.yield()
   end
 end
 
+function Task:prefix()
+  local plugin = "[" .. self.plugin.name .. "] "
+  local task = string.rep(" ", 20 - #(self.name .. self.plugin.name)) .. self.name
+
+  return colors and Terminal.magenta(plugin) .. Terminal.cyan(task) .. Terminal.bright_black(" | ")
+    or plugin .. " " .. task .. " | "
+end
+
 function Task:wait()
   while self:is_running() do
     vim.wait(10)
diff --git a/lua/lazy/terminal.lua b/lua/lazy/terminal.lua
new file mode 100644
index 0000000..a1bcb84
--- /dev/null
+++ b/lua/lazy/terminal.lua
@@ -0,0 +1,71 @@
+---@class Ansi: table<string, fun(string):string>
+local M = {}
+
+M.colors = {
+  reset = "\27[0m",
+  black = "\27[30m",
+  red = "\27[31m",
+  green = "\27[32m",
+  yellow = "\27[33m",
+  blue = "\27[34m",
+  magenta = "\27[35m",
+  cyan = "\27[36m",
+  white = "\27[37m",
+  bright_black = "\27[90m",
+  bright_red = "\27[91m",
+  bright_green = "\27[92m",
+  bright_yellow = "\27[93m",
+  bright_blue = "\27[94m",
+  bright_magenta = "\27[95m",
+  bright_cyan = "\27[96m",
+  bright_white = "\27[97m",
+}
+
+function M.color(text, color)
+  return M.colors[color] .. text .. M.colors.reset
+end
+
+-- stylua: ignore start
+function M.black(text) return M.color(text, "black") end
+function M.red(text) return M.color(text, "red") end
+function M.green(text) return M.color(text, "green") end
+function M.yellow(text) return M.color(text, "yellow") end
+function M.blue(text) return M.color(text, "blue") end
+function M.magenta(text) return M.color(text, "magenta") end
+function M.cyan(text) return M.color(text, "cyan") end
+function M.white(text) return M.color(text, "white") end
+function M.bright_black(text) return M.color(text, "bright_black") end
+function M.bright_red(text) return M.color(text, "bright_red") end
+function M.bright_green(text) return M.color(text, "bright_green") end
+function M.bright_yellow(text) return M.color(text, "bright_yellow") end
+function M.bright_blue(text) return M.color(text, "bright_blue") end
+function M.bright_magenta(text) return M.color(text, "bright_magenta") end
+function M.bright_cyan(text) return M.color(text, "bright_cyan") end
+function M.bright_white(text) return M.color(text, "bright_white") end
+-- stylua: ignore end
+
+---@param data string
+---@param prefix string
+function M.prefix(data, prefix)
+  -- Normalize Windows-style newlines to simple newlines
+  data = data:gsub("\r\n", "\n")
+
+  -- Handle prefix for the first line, if data starts immediately
+  data = prefix .. data
+
+  -- Prefix new lines ensuring not to double prefix if a line starts with \r
+  data = data:gsub("(\n)([^\r])", "%1" .. prefix .. "%2")
+
+  -- Handle carriage returns properly to avoid double prefixing
+  -- Replace any \r not followed by \n with \r, then add a prefix only if the following character isn't the start of our prefix
+  data = data:gsub("\r([^\n])", function(nextChar)
+    if nextChar:sub(1, #prefix) == prefix then
+      return "\r" .. nextChar
+    else
+      return "\r" .. prefix .. nextChar
+    end
+  end)
+  return data
+end
+
+return M
diff --git a/tests/busted.lua b/tests/busted.lua
index 146ab96..8fd4f28 100755
--- a/tests/busted.lua
+++ b/tests/busted.lua
@@ -6,22 +6,22 @@ for _, name in ipairs({ "config", "data", "state", "cache" }) do
   vim.env[("XDG_%s_HOME"):format(name:upper())] = root .. "/" .. name
 end
 
--- -- Bootstrap lazy.nvim
--- local lazypath = vim.fn.stdpath("data") .. "/lazy/lazy.nvim"
--- if not (vim.uv or vim.loop).fs_stat(lazypath) then
---   local lazyrepo = "https://github.com/folke/lazy.nvim.git"
---   vim.fn.system({ "git", "clone", "--filter=blob:none", "--branch=stable", lazyrepo, lazypath })
--- end
--- vim.opt.rtp:prepend(lazypath)
 vim.opt.rtp:prepend(".")
 
 vim.o.loadplugins = true -- enable since nvim -l disables plugins
 
 -- Setup lazy.nvim
 require("lazy").setup({
-  "lunarmodules/busted", -- add busted
+  spec = {
+    "lunarmodules/busted", -- add busted
+  },
+  rocks = { hererocks = true },
 })
 
+local Config = require("lazy.core.config")
+-- disable termnial output for the tests
+Config.options.headless = {}
+
 -- run busted
 return pcall(require("busted.runner"), {
   standalone = false,