feat(rst): support treesitter to get headings

README.md

@@ -30,7 +30,7 @@ You can setup the extension by adding the following to your config:
 require('telescope').load_extension('heading')
 ```
 
-Use treesitter (which only supports Markdown) to query headings. Setup with `telescope.setup` **before** `load_extension`:
+Use treesitter (which only supports Markdown/ReStructuredText) to query headings. Setup with `telescope.setup` **before** `load_extension`:
 ```lua
 local telescope = require('telescope')
 telescope.setup({

lua/telescope/_extensions/heading.lua

@@ -28,10 +28,18 @@ local function filetype()
     return ft
 end
 
+local function support_treesitter(ft)
+    if ft == 'markdown' or ft == 'rst' then
+        return true
+    end
+    return false
+end
+
 local function get_headings()
+    local ft = filetype()
     local mod_path = string.format(
         'telescope._extensions.heading.format.%s',
-        filetype()
+        ft
     )
     local ok, mod = pcall(require, mod_path)
     if not ok then
@@ -46,7 +54,12 @@ local function get_headings()
     local index, total = 1, vim.fn.line('$')
     local bufnr = vim.api.nvim_get_current_buf()
     local filepath = vim.api.nvim_buf_get_name(bufnr)
-    local headings = mod.get_headings(filepath, index, total)
+    local headings = {} -- luacheck: ignore
+    if heading_config.treesitter and support_treesitter(ft) then
+        headings = mod.ts_get_headings(filepath, bufnr)
+    else
+        headings = mod.get_headings(filepath, index, total)
+    end
     return headings
 end
 

lua/telescope/_extensions/heading/format/markdown.lua

@@ -1,12 +1,7 @@
-local heading_config = require('telescope._extensions.heading.config')
-
 local Markdown = {}
 
 function Markdown.get_headings(filepath, start, total)
     local headings = {}
-    if heading_config.treesitter then
-        return Markdown.ts_get_headings(filepath)
-    end
     local index = start
     local matches = {
         '# ',
@@ -47,13 +42,12 @@ function Markdown.get_headings(filepath, start, total)
     return headings
 end
 
-function Markdown.ts_get_headings(filepath)
+function Markdown.ts_get_headings(filepath, bufnr)
     local ts = vim.treesitter
     local query = [[
     (atx_heading) @heading
     ]]
     local parsed_query = ts.parse_query('markdown', query)
-    local bufnr = vim.api.nvim_get_current_buf() or 0
     local parser = ts.get_parser(bufnr, 'markdown')
     local root = parser:parse()[1]:root()
     local start_row, _, end_row, _ = root:range()

lua/telescope/_extensions/heading/format/rst.lua

@@ -86,4 +86,27 @@ function ReStructuredText.get_headings(filepath, start, total)
     return headings
 end
 
+function ReStructuredText.ts_get_headings(filepath, bufnr)
+    local ts = vim.treesitter
+    local query = [[
+    (section) (title) @section_title
+    ]]
+    local parsed_query = ts.parse_query('rst', query)
+    local parser = ts.get_parser(bufnr, 'rst')
+    local root = parser:parse()[1]:root()
+    local start_row, _, end_row, _ = root:range()
+
+    local headings = {}
+    for _, node in parsed_query:iter_captures(root, bufnr, start_row, end_row) do
+        local row, _ = node:range()
+        local line = vim.fn.getline(row + 1)
+        table.insert(headings, {
+            heading = vim.trim(line),
+            line = row + 1,
+            path = filepath,
+        })
+    end
+    return headings
+end
+
 return ReStructuredText