Added syntax highlighting for cell separators and magic commands

README.md

@@ -184,6 +184,17 @@ require("jupynium").setup({
     use_default_keybindings = true,
   },
 
+  syntax_highlight = {
+    enable = true,
+
+    highlight_groups = {
+      -- Set to nil to disable each entry
+      code_cell_separator = "Folded",
+      markdown_cell_separator = "Pmenu",
+      magic_command = "Keyword",
+    },
+  },
+
   -- Dim all cells except the current one
   -- Related command :JupyniumShortsightedToggle
   shortsighted = false,

lua/jupynium/cells.lua

@@ -1,40 +1,63 @@
+local utils = require "jupynium.utils"
+
 local M = {}
 
-local string_starts_with = function(str, start)
-  if str == nil or start == nil then
-    return false
+--- Get the line type (cell separator, magic commands, empty, others)
+---@param line string | number 1-indexed
+---@return string "cell separator: markdown" | "cell separator: markdown (jupytext)" | "cell separator: code" | "magic commands" | "empty" | "others"
+function M.line_type(line)
+  if type(line) == "number" then
+    line = vim.api.nvim_buf_get_lines(0, line - 1, line, false)[1]
   end
-  return str:sub(1, #start) == start
-end
 
-local is_line_separator = function(row)
-  local line = vim.api.nvim_buf_get_lines(0, row - 1, row, false)[1]
   if
-    string_starts_with(line, "# %%")
-    or string_starts_with(line, "# %%%")
-    or string_starts_with(line, "# %% [md]")
-    or string_starts_with(line, "# %% [markdown]")
-    or string_starts_with(line, '"""%%')
-    or string_starts_with(line, "'''%%")
-    or string_starts_with(line, '%%"""')
-    or string_starts_with(line, "%%'''")
+    utils.string_begins_with(line, "# %%%")
+    or utils.string_begins_with(line, '"""%%')
+    or utils.string_begins_with(line, "'''%%")
+  then
+    return "cell separator: markdown"
+  elseif utils.string_begins_with(line, "# %% [md]") or utils.string_begins_with(line, "# %% [markdown]") then
+    return "cell separator: markdown (jupytext)"
+  elseif
+    utils.string_begins_with(line, "# %%")
+    or utils.string_begins_with(line, '%%"""')
+    or utils.string_begins_with(line, "%%'''")
   then
+    return "cell separator: code"
+  elseif utils.string_begins_with(line, "# %") then
+    return "magic command"
+  elseif vim.fn.trim(line) == "" then
+    return "empty"
+  end
+
+  return "others" -- code
+end
+
+--- Check if the line is a cell separator
+---@param line string | number 1-indexed
+---@return boolean
+function M.is_line_separator(line)
+  local line_type = M.line_type(line)
+  if utils.string_begins_with(line_type, "cell separator:") then
     return true
   end
+
   return false
 end
 
+--- Get the current cell separator row
+---@param row number | nil 1-indexed
+---@return number | nil row 1-indexed
 function M.current_cell_separator(row)
   row = row or vim.api.nvim_win_get_cursor(0)[1]
-  local found_separator = 0
-  if is_line_separator(row) then
+  if M.is_line_separator(row) then
     return row
   end
 
   row = row - 1
 
   while row > 0 do
-    if is_line_separator(row) then
+    if M.is_line_separator(row) then
       return row
     end
     row = row - 1
@@ -43,18 +66,21 @@ function M.current_cell_separator(row)
   return nil
 end
 
+--- Get the previous cell separator row
+---@param row number | nil 1-indexed
+---@return number | nil row 1-indexed
 function M.previous_cell_separator(row)
   row = row or vim.api.nvim_win_get_cursor(0)[1]
-  local found_separator = 0
-  if is_line_separator(row) then
+  local found_separator
+  if M.is_line_separator(row) then
     found_separator = row
   end
 
   row = row - 1
 
   while row > 0 do
-    if is_line_separator(row) then
-      if found_separator > 0 then
+    if M.is_line_separator(row) then
+      if found_separator ~= nil then
         return row
       else
         found_separator = row
@@ -63,20 +89,24 @@ function M.previous_cell_separator(row)
     row = row - 1
   end
 
-  if found_separator > 0 then
+  if found_separator ~= nil then
     return found_separator
   end
   return nil
 end
 
+--- Get the next cell separator row
+---@param row number | nil 1-indexed
+---@return number | nil row 1-indexed
 function M.next_cell_separator(row)
   row = row or vim.api.nvim_win_get_cursor(0)[1]
+
   row = row + 1
 
   local num_lines = vim.api.nvim_buf_line_count(0)
 
   while row <= num_lines do
-    if is_line_separator(row) then
+    if M.is_line_separator(row) then
       return row
     end
     row = row + 1

lua/jupynium/highlighter.lua

@@ -0,0 +1,178 @@
+-- Code mostly based on koenverburg/peepsight.nvim
+local utils = require "jupynium.utils"
+local cells = require "jupynium.cells"
+local M = {}
+
+M.options = {
+  enable = true, -- separate from shortsighted.enable. Only for highlighting.
+  highlight_groups = {},
+
+  shortsighted = {
+    enable = true,
+    highlight_groups = {
+      dim = "Comment",
+    },
+  },
+}
+
+function M.is_enabled()
+  return M.options.enable or M.options.shortsighted.enable
+end
+
+function M.set_autocmd()
+  local augroup = vim.api.nvim_create_augroup("jupynium-highlighter", {})
+  vim.api.nvim_create_autocmd({ "BufWinEnter", "BufWritePost", "CursorMoved", "CursorMovedI", "WinScrolled" }, {
+    pattern = "*.ju.*",
+    callback = M.run,
+    group = augroup,
+  })
+end
+
+local ns_highlight = vim.api.nvim_create_namespace "jupynium-highlighter"
+local ns_shortsighted = vim.api.nvim_create_namespace "jupynium-shortsighted"
+
+--- Set highlight group for a line
+---@param buffer number
+---@param line_number number 0-indexed
+---@param hl_group string
+function M.set_line_hlgroup(buffer, namespace, line_number, hl_group)
+  pcall(vim.api.nvim_buf_set_extmark, buffer, namespace, line_number, 0, {
+    end_line = line_number + 1,
+    end_col = 0,
+    hl_group = hl_group,
+    hl_eol = true,
+    priority = 10000,
+  })
+end
+
+function M.clear_namespace(namespace)
+  vim.api.nvim_buf_clear_namespace(0, namespace, 0, -1)
+end
+
+function M.enable()
+  M.options.enable = true
+
+  if string.find(vim.fn.expand "%", ".ju.") then
+    M.update()
+  end
+  M.set_autocmd()
+end
+
+function M.disable()
+  M.options.enable = false
+
+  M.clear_namespace(ns_highlight)
+end
+
+function M.toggle()
+  if M.options.enable then
+    M.disable()
+  else
+    M.enable()
+  end
+end
+
+function M.run()
+  if M.is_enabled() then
+    M.clear_namespace(ns_highlight)
+    M.clear_namespace(ns_shortsighted)
+
+    M.update()
+  end
+end
+
+function M.shortsighted_enable()
+  M.options.shortsighted.enable = true
+
+  if string.find(vim.fn.expand "%", ".ju.") then
+    M.update()
+  end
+  M.set_autocmd()
+end
+
+function M.shortsighted_disable()
+  M.options.shortsighted.enable = false
+
+  M.clear_namespace(ns_shortsighted)
+end
+
+function M.shortsighted_toggle()
+  if M.options.shortsighted.enable then
+    M.shortsighted_disable()
+  else
+    M.shortsighted_enable()
+  end
+end
+
+function M.update()
+  if not M.is_enabled() then
+    return
+  end
+
+  local end_of_file = vim.fn.line "$"
+  local lines = vim.api.nvim_buf_get_lines(0, 0, -1, false)
+  local line_types = {}
+  for i, line in ipairs(lines) do
+    line_types[i] = cells.line_type(line)
+  end
+
+  if M.options.enable then
+    for i, line_type in ipairs(line_types) do
+      if line_type == "cell separator: code" then
+        M.set_line_hlgroup(0, ns_highlight, i - 1, M.options.highlight_groups.code_cell_separator)
+      elseif utils.string_begins_with(line_type, "cell separator: markdown") then
+        M.set_line_hlgroup(0, ns_highlight, i - 1, M.options.highlight_groups.markdown_cell_separator)
+      elseif line_type == "magic command" then
+        M.set_line_hlgroup(0, ns_highlight, i - 1, M.options.highlight_groups.magic_command)
+      end
+    end
+  end
+
+  if M.options.shortsighted.enable then
+    local current_row = cells.current_cell_separator()
+    local next_row = cells.next_cell_separator()
+
+    if current_row == nil and next_row == nil then
+      return
+    end
+
+    if current_row ~= nil then
+      -- Dim above cell range
+      -- Exclude current cell separator
+      for i = 1, current_row - 1 do
+        if line_types[i] ~= "empty" then
+          M.set_line_hlgroup(0, ns_shortsighted, i - 1, M.options.shortsighted.highlight_groups.dim)
+        end
+      end
+    end
+
+    if next_row ~= nil then
+      -- Dim below cell range
+      for j = next_row, end_of_file do
+        if not line_types[j] ~= "empty" then
+          M.set_line_hlgroup(0, ns_shortsighted, j - 1, M.options.shortsighted.highlight_groups.dim)
+        end
+      end
+    end
+  end
+end
+
+function M.add_commands()
+  vim.api.nvim_create_user_command(
+    "JupyniumShortsightedToggle",
+    "lua require('jupynium.highlighter').shortsighted_toggle()",
+    {}
+  )
+  vim.api.nvim_create_user_command(
+    "JupyniumShortsightedEnable",
+    "lua require('jupynium.highlighter').shortsighted_enable()",
+    {}
+  )
+  vim.api.nvim_create_user_command(
+    "JupyniumShortsightedDisable",
+    "lua require('jupynium.highlighter').shortsighted_disable()",
+    {}
+  )
+end
+
+return M

lua/jupynium/init.lua

@@ -1,7 +1,7 @@
 local M = {}
 
 local textobj = require "jupynium.textobj"
-local shortsighted = require "jupynium.shortsighted"
+local highlighter = require "jupynium.highlighter"
 local server = require "jupynium.server"
 local options = require "jupynium.options"
 
@@ -83,12 +83,19 @@ function M.setup(opts)
     textobj.default_keybindings(augroup)
   end
 
+  highlighter.options.highlight_groups = options.opts.syntax_highlight.highlight_groups
+  if options.opts.syntax_highlight.enable then
+    highlighter.enable()
+  else
+    highlighter.disable()
+  end
+
   if options.opts.shortsighted then
-    shortsighted.enable()
+    highlighter.shortsighted_enable()
   else
-    shortsighted.disable()
+    highlighter.shortsighted_disable()
   end
-  shortsighted.add_commands()
+  highlighter.add_commands()
 
   vim.g.__jupynium_setup_completed = true
 end

lua/jupynium/options.lua

@@ -84,6 +84,17 @@ M.default_opts = {
     use_default_keybindings = true,
   },
 
+  syntax_highlight = {
+    enable = true,
+
+    highlight_groups = {
+      -- Set to nil to disable each entry
+      code_cell_separator = "Folded",
+      markdown_cell_separator = "Pmenu",
+      magic_command = "Keyword",
+    },
+  },
+
   -- Dim all cells except the current one
   -- Related command :JupyniumShortsightedToggle
   shortsighted = false,