From ee2374ecb5125b580e0158d71957f5dda9a49d5a Mon Sep 17 00:00:00 2001
From: Tristan Knight <tris203@gmail.com>
Date: Fri, 29 Dec 2023 22:46:37 +0000
Subject: [PATCH] fix: lazy mapped keys with functions (#58)

---
 lua/hawtkeys/ts.lua        | 11 ++++++++++-
 lua/hawtkeys/ui.lua        |  7 ++++---
 lua/hawtkeys/utils.lua     |  7 +++++++
 tests/hawtkeys/ts_spec.lua | 30 ++++++++++++++++++++++++++++--
 4 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/lua/hawtkeys/ts.lua b/lua/hawtkeys/ts.lua
index 94e8bd7..96640af 100644
--- a/lua/hawtkeys/ts.lua
+++ b/lua/hawtkeys/ts.lua
@@ -376,9 +376,18 @@ local function get_keymaps_from_lazy()
             for _, v in ipairs(lazy) do
                 if v and v._ and v._.handlers and v._.handlers.keys then
                     for _, key in pairs(v._.handlers.keys) do
+                        if type(key.rhs) == "table" then
+                            key.rhs = tostring(key.rhs)
+                        elseif type(key.rhs) == "function" then
+                            local debugInfo =
+                                debug.getinfo(key.rhs --[[@as fun()]], "S")
+                            key.rhs = utils.reduceHome(debugInfo.short_src)
+                                .. ":"
+                                .. debugInfo.linedefined
+                        end
                         local map = {
                             lhs = key.lhs,
-                            rhs = key.rhs,
+                            rhs = tostring(key.rhs),
                             mode = key.mode,
                             from_file = "Lazy Init:" .. tostring(v[1]),
                         }
diff --git a/lua/hawtkeys/ui.lua b/lua/hawtkeys/ui.lua
index 33f86a1..8baeb3b 100644
--- a/lua/hawtkeys/ui.lua
+++ b/lua/hawtkeys/ui.lua
@@ -3,6 +3,7 @@ local M = {}
 local Hawtkeys = require("hawtkeys.score")
 local ShowAll = require("hawtkeys.show_all")
 local showDuplicates = require("hawtkeys.duplicates")
+local utils = require("hawtkeys.utils")
 
 local Namespace = vim.api.nvim_create_namespace("hawtkeys")
 
@@ -263,7 +264,7 @@ M.show_all = function()
     local all = ShowAll.show_all()
     local pattern = "%s (%s) - %s"
     for i, data in ipairs(all) do
-        local filename = data.from_file:gsub(vim.env.HOME, "~")
+        local filename = utils.reduceHome(data.from_file)
         local line = pattern:format(data.lhs, data.mode, filename)
 
         local offset_mode = #data.lhs + 2
@@ -319,8 +320,8 @@ M.show_dupes = function()
     local dupes = showDuplicates.show_duplicates()
     local pattern = "%s : %s"
     for i, data in ipairs(dupes) do
-        local filename1 = data.file1:gsub(vim.env.HOME, "~")
-        local filename2 = data.file2:gsub(vim.env.HOME, "~")
+        local filename1 = utils.reduceHome(data.file1)
+        local filename2 = utils.reduceHome(data.file2)
         local line = pattern:format(filename1, filename2)
 
         local l2 = data.key
diff --git a/lua/hawtkeys/utils.lua b/lua/hawtkeys/utils.lua
index c2b53b1..bb8bfc3 100644
--- a/lua/hawtkeys/utils.lua
+++ b/lua/hawtkeys/utils.lua
@@ -80,4 +80,11 @@ function M.find_duplicates(keymaps)
     return duplicates
 end
 
+---@param path string
+---@return string
+function M.reduceHome(path)
+    local reduced = path:gsub(vim.env.HOME, "~")
+    return reduced
+end
+
 return M
diff --git a/tests/hawtkeys/ts_spec.lua b/tests/hawtkeys/ts_spec.lua
index 66ef80a..bfb2044 100644
--- a/tests/hawtkeys/ts_spec.lua
+++ b/tests/hawtkeys/ts_spec.lua
@@ -174,6 +174,8 @@ describe("Lazy Managed Plugins", function()
                 package.loaded[modname] = nil
             end
         end
+    end)
+    it("extract keys set in a Lazy init", function()
         local lazy = require("lazy")
         lazy.setup({
             "ellisonleao/nvim-plugin-template",
@@ -185,8 +187,6 @@ describe("Lazy Managed Plugins", function()
                 },
             },
         })
-    end)
-    it("extract keys set in a Lazy init", function()
         local lazyKeymaps = ts.get_keymaps_from_lazy()
         eq("n", lazyKeymaps[1].mode)
         eq("<leader>1", lazyKeymaps[1].lhs)
@@ -196,4 +196,30 @@ describe("Lazy Managed Plugins", function()
             lazyKeymaps[1].from_file
         )
     end)
+    it(
+        "extract keys where the lazy key setting is a function and returns as string",
+        function()
+            local lazy = require("lazy")
+            lazy.setup({
+                "ellisonleao/nvim-plugin-template",
+                keys = {
+                    {
+                        "<leader>1",
+                        function()
+                            print(1)
+                        end,
+                        desc = "Test Lazy Print 1",
+                    },
+                },
+            })
+            local lazyKeymaps = ts.get_keymaps_from_lazy()
+            eq("n", lazyKeymaps[1].mode)
+            eq("<leader>1", lazyKeymaps[1].lhs)
+            eq("string", type(lazyKeymaps[1].rhs))
+            eq(
+                "Lazy Init:ellisonleao/nvim-plugin-template",
+                lazyKeymaps[1].from_file
+            )
+        end
+    )
 end)