Merge pull request #228 from stevearc/stevearc-refactor-hl

refactor: change hl group calculation to be more customizable

README.md

@@ -424,6 +424,11 @@ require("aerial").setup({
     whitespace = "  ",
   },
 
+  -- Set this function to override the highlight groups for certain symbols
+  get_highlight = function(symbol, is_icon)
+    -- return "MyHighlight" .. symbol.kind
+  end,
+
   -- Options for opening aerial in a floating win
   float = {
     -- Controls border appearance. Passed to nvim_open_win

doc/aerial.txt

@@ -262,6 +262,11 @@ OPTIONS                                                           *aerial-option
         whitespace = "  ",
       },
 
+      -- Set this function to override the highlight groups for certain symbols
+      get_highlight = function(symbol, is_icon)
+        -- return "MyHighlight" .. symbol.kind
+      end,
+
       -- Options for opening aerial in a floating win
       float = {
         -- Controls border appearance. Passed to nvim_open_win

lua/aerial/config.lua

@@ -248,6 +248,11 @@ local default_options = {
     whitespace = "  ",
   },
 
+  -- Set this function to override the highlight groups for certain symbols
+  get_highlight = function(symbol, is_icon)
+    -- return "MyHighlight" .. symbol.kind
+  end,
+
   -- Options for opening aerial in a floating win
   float = {
     -- Controls border appearance. Passed to nvim_open_win

lua/aerial/data.lua

@@ -1,6 +1,8 @@
 local util = require("aerial.util")
 local config = require("aerial.config")
 
+---@alias aerial.Scope "private"|"public"|"protected"
+
 ---@class aerial.Range
 ---@field lnum integer
 ---@field end_lnum integer
@@ -11,6 +13,7 @@ local config = require("aerial.config")
 ---@field kind string
 ---@field name string
 ---@field level integer
+---@field scope nil|aerial.Scope
 ---@field parent? aerial.Symbol
 ---@field selection_range? aerial.Range
 ---@field children? aerial.Symbol[]

lua/aerial/highlight.lua

@@ -1,3 +1,4 @@
+local config = require("aerial.config")
 local M = {}
 
 ---@param group1 string
@@ -6,11 +7,61 @@ local function link(group1, group2)
   vim.api.nvim_set_hl(0, group1, { link = group2, default = true })
 end
 
+local symbol_kinds = {
+  "Array",
+  "Boolean",
+  "Class",
+  "Constant",
+  "Constructor",
+  "Enum",
+  "EnumMember",
+  "Event",
+  "Field",
+  "File",
+  "Function",
+  "Interface",
+  "Key",
+  "Method",
+  "Module",
+  "Namespace",
+  "Null",
+  "Number",
+  "Object",
+  "Operator",
+  "Package",
+  "Property",
+  "String",
+  "Struct",
+  "TypeParameter",
+  "Variable",
+}
+
+---@param symbol aerial.Symbol
+---@param is_icon boolean
+---@return nil|string
+M.get_highlight = function(symbol, is_icon)
+  local hl_group = config.get_highlight(symbol, is_icon)
+  if hl_group then
+    return hl_group
+  end
+
+  -- Default functionality will set the highlight for private and protected symbol names
+  if symbol.scope and not is_icon and symbol.scope ~= "public" then
+    return string.format("Aerial%s", symbol.scope:gsub("^%l", string.upper))
+  end
+
+  return string.format("Aerial%s%s", symbol.kind, is_icon and "Icon" or "")
+end
+
 M.create_highlight_groups = function()
   -- The line that shows where your cursor(s) are
   link("AerialLine", "QuickFixLine")
   link("AerialLineNC", "AerialLine")
 
+  -- Highlight groups for private and protected functions/fields/etc
+  link("AerialPrivate", "Comment")
+  link("AerialProtected", "Comment")
+
   -- Use Comment colors for AerialGuide, while stripping bold/italic/etc
   local comment_defn = vim.api.nvim_get_hl_by_name("Comment", true)
   -- The guides when show_guide = true
@@ -27,7 +78,7 @@ M.create_highlight_groups = function()
   end
 
   -- The name of the symbol
-  for _, symbol_kind in ipairs(M.identifiers) do
+  for _, symbol_kind in ipairs(symbol_kinds) do
     link(string.format("Aerial%s", symbol_kind), "NONE")
   end
 
@@ -60,33 +111,4 @@ M.create_highlight_groups = function()
   link("AerialVariableIcon", "Identifier")
 end
 
-M.identifiers = {
-  "Array",
-  "Boolean",
-  "Class",
-  "Constant",
-  "Constructor",
-  "Enum",
-  "EnumMember",
-  "Event",
-  "Field",
-  "File",
-  "Function",
-  "Interface",
-  "Key",
-  "Method",
-  "Module",
-  "Namespace",
-  "Null",
-  "Number",
-  "Object",
-  "Operator",
-  "Package",
-  "Property",
-  "String",
-  "Struct",
-  "TypeParameter",
-  "Variable",
-}
-
 return M

lua/aerial/render.lua

@@ -1,6 +1,7 @@
 local backends = require("aerial.backends")
 local config = require("aerial.config")
 local data = require("aerial.data")
+local highlight = require("aerial.highlight")
 local layout = require("aerial.layout")
 local loading = require("aerial.loading")
 local util = require("aerial.util")
@@ -150,18 +151,24 @@ M.update_aerial_buffer = function(buf)
     end
     local text = util.remove_newlines(string.format("%s%s %s", spacing, kind, item.name))
     local text_cols = vim.api.nvim_strwidth(text)
-    table.insert(highlights, {
-      group = "Aerial" .. item.kind .. "Icon",
-      row = row,
-      col_start = string_len[spacing],
-      col_end = string_len[spacing] + string_len[kind],
-    })
-    table.insert(highlights, {
-      group = "Aerial" .. item.kind,
-      row = row,
-      col_start = string_len[spacing] + string_len[kind],
-      col_end = -1,
-    })
+    local icon_hl = highlight.get_highlight(item, true)
+    if icon_hl then
+      table.insert(highlights, {
+        group = icon_hl,
+        row = row,
+        col_start = string_len[spacing],
+        col_end = string_len[spacing] + string_len[kind],
+      })
+    end
+    local name_hl = highlight.get_highlight(item, false)
+    if name_hl then
+      table.insert(highlights, {
+        group = name_hl,
+        row = row,
+        col_start = string_len[spacing] + string_len[kind],
+        col_end = -1,
+      })
+    end
     max_len = math.max(max_len, text_cols)
     table.insert(lines, text)
     row = row + 1

lua/lualine/components/aerial.lua

@@ -43,7 +43,6 @@
 --  my_method_name
 
 local M = require("lualine.component"):extend()
-local aerial = require("aerial")
 local utils = require("lualine.utils.utils")
 
 local default_options = {
@@ -54,21 +53,6 @@ local default_options = {
   exact = true,
 }
 
-function M:color_for_lualine()
-  local identifiers = require("aerial.highlight").identifiers
-  self.highlight_groups = {}
-  for _, symbol_kind in ipairs(identifiers) do
-    local hl = "Aerial" .. symbol_kind
-    local hl_icon = "Aerial" .. symbol_kind .. "Icon"
-    local color = { fg = utils.extract_highlight_colors(hl, "fg") }
-    local color_icon = { fg = utils.extract_highlight_colors(hl_icon, "fg") }
-    self.highlight_groups[symbol_kind] = {
-      icon = self:create_hl(color_icon, symbol_kind .. "Icon"),
-      text = self:create_hl(color, symbol_kind),
-    }
-  end
-end
-
 function M:format_status(symbols, depth, separator, icons_enabled, colored)
   local parts = {}
   depth = depth or #symbols
@@ -80,15 +64,9 @@ function M:format_status(symbols, depth, separator, icons_enabled, colored)
   end
 
   for _, symbol in ipairs(symbols) do
-    local name = utils.stl_escape(symbol.name)
-    if colored then
-      name = self:format_hl(self.highlight_groups[symbol.kind].text) .. name
-    end
+    local name = self:color_text(utils.stl_escape(symbol.name), symbol, false)
     if icons_enabled then
-      local icon = symbol.icon
-      if colored then
-        icon = self:format_hl(self.highlight_groups[symbol.kind].icon) .. icon
-      end
+      local icon = self:color_text(symbol.icon, symbol, true)
       table.insert(parts, string.format("%s %s", icon, name))
     else
       table.insert(parts, name)
@@ -104,11 +82,7 @@ function M:init(options)
   if self.options.colored == nil then
     self.options.colored = true
   end
-  if self.options.colored then
-    aerial.sync_load()
-    require("aerial.highlight").create_highlight_groups()
-    self:color_for_lualine()
-  end
+  self.highlight_groups = {}
   self.get_status = self.get_status_normal
 
   if self.options.dense then
@@ -116,12 +90,34 @@ function M:init(options)
   end
 end
 
+---@param text string
+---@param symbol aerial.Symbol
+---@param is_icon boolean
+---@return string
+function M:color_text(text, symbol, is_icon)
+  if not self.options.colored then
+    return text
+  end
+  local hl_group = require("aerial.highlight").get_highlight(symbol, is_icon)
+  if hl_group then
+    local lualine_hl_group = self.highlight_groups[hl_group]
+    if not lualine_hl_group then
+      lualine_hl_group =
+        self:create_hl({ fg = utils.extract_highlight_colors(hl_group, "fg") }, "LL" .. hl_group)
+      self.highlight_groups[hl_group] = lualine_hl_group
+    end
+    return self:format_hl(lualine_hl_group) .. text
+  else
+    return text
+  end
+end
+
 function M:update_status()
   return self:get_status()
 end
 
 function M:get_status_normal()
-  local symbols = aerial.get_location(self.options.exact)
+  local symbols = require("aerial").get_location(self.options.exact)
   local status = self:format_status(
     symbols,
     self.options.depth,
@@ -133,7 +129,7 @@ function M:get_status_normal()
 end
 
 function M:get_status_dense()
-  local symbols = aerial.get_location(self.options.exact)
+  local symbols = require("aerial").get_location(self.options.exact)
   local status = self:format_status(
     symbols,
     self.options.depth,
@@ -146,10 +142,7 @@ function M:get_status_dense()
 
   if self.options.icons_enabled and not vim.tbl_isempty(symbols) then
     local symbol = symbols[#symbols]
-    local icon = symbol.icon
-    if self.options.colored then
-      icon = self:format_hl(self.highlight_groups[symbol.kind].icon) .. icon
-    end
+    local icon = self:color_text(symbol.icon, symbol, true)
     status = string.format("%s %s", icon, status)
   end
   return status

lua/telescope/_extensions/aerial.lua

@@ -18,6 +18,7 @@ local function aerial_picker(opts)
   local backends = require("aerial.backends")
   local config = require("aerial.config")
   local data = require("aerial.data")
+  local highlight = require("aerial.highlight")
   local util = require("aerial.util")
 
   local bufnr = vim.api.nvim_get_current_buf()
@@ -51,9 +52,11 @@ local function aerial_picker(opts)
     local icon = config.get_icon(bufnr, item.kind)
     local text = vim.api.nvim_buf_get_lines(bufnr, item.lnum - 1, item.lnum, false)[1] or ""
     text = vim.trim(text)
+    local icon_hl = highlight.get_highlight(item, true) or "NONE"
+    local name_hl = highlight.get_highlight(item, false) or "NONE"
     local columns = {
-      { icon, "Aerial" .. item.kind .. "Icon" },
-      { entry.name, "Aerial" .. item.kind },
+      { icon, icon_hl },
+      { entry.name, name_hl },
       text,
     }
     return displayer(columns)