From f685f2d5da05ecd19e83fa5a3c655147d0cad9d3 Mon Sep 17 00:00:00 2001
From: David Zhang <crispgm@gmail.com>
Date: Wed, 13 Apr 2022 14:59:36 +0800
Subject: [PATCH] feat(rst): support treesitter to get headings

---
 README.md                                     |  2 +-
 lua/telescope/_extensions/heading.lua         | 17 ++++++++++++--
 .../_extensions/heading/format/markdown.lua   |  8 +------
 .../_extensions/heading/format/rst.lua        | 23 +++++++++++++++++++
 4 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index c5fa64a..eefc5ea 100644
--- a/README.md
+++ b/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({
diff --git a/lua/telescope/_extensions/heading.lua b/lua/telescope/_extensions/heading.lua
index 51c726d..82a4b26 100644
--- a/lua/telescope/_extensions/heading.lua
+++ b/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
 
diff --git a/lua/telescope/_extensions/heading/format/markdown.lua b/lua/telescope/_extensions/heading/format/markdown.lua
index 567ef42..cfdd0af 100644
--- a/lua/telescope/_extensions/heading/format/markdown.lua
+++ b/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()
diff --git a/lua/telescope/_extensions/heading/format/rst.lua b/lua/telescope/_extensions/heading/format/rst.lua
index 5cdca47..880d8d9 100644
--- a/lua/telescope/_extensions/heading/format/rst.lua
+++ b/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