From 5c420844227c75390cc9fdf6047bfc49466169d9 Mon Sep 17 00:00:00 2001
From: champignoom <66909116+champignoom@users.noreply.github.com>
Date: Tue, 28 Nov 2023 02:56:30 +0800
Subject: [PATCH] feat(toc): add toc item filter (#1195)

Co-authored-by: vhyrro <76052559+vhyrro@users.noreply.github.com>
---
 lua/neorg/modules/core/qol/toc/module.lua | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/lua/neorg/modules/core/qol/toc/module.lua b/lua/neorg/modules/core/qol/toc/module.lua
index de29d788c..47d8a2ddf 100644
--- a/lua/neorg/modules/core/qol/toc/module.lua
+++ b/lua/neorg/modules/core/qol/toc/module.lua
@@ -51,6 +51,10 @@ module.config.public = {
     -- If `true`, the width of the Table of Contents window will automatically
     -- fit its longest line
     fit_width = true,
+
+    -- A function that takes node type as argument and returns true if this
+    -- type of node should appear in the Table of Contents
+    toc_filter = function(node_type) return node_type:match("^heading") end,
 }
 
 module.public = {
@@ -149,6 +153,7 @@ module.public = {
 
         local prefix, title
         local extmarks = {}
+        local toc_filter = module.config.public.toc_filter
 
         local success = module.required["core.integrations.treesitter"].execute_query(
             [[
@@ -163,7 +168,7 @@ module.public = {
                 local capture = query.captures[id]
 
                 if capture == "prefix" then
-                    if node:type():match("_prefix$") then
+                    if node:type():match("_prefix$") and (type(toc_filter) ~= 'function' or toc_filter(node:type())) then
                         prefix = node
                     else
                         prefix = nil