Full Code of GCBallesteros/NotebookNavigator.nvim for AI

Repository: GCBallesteros/NotebookNavigator.nvim
Branch: main
Commit: 20cb6f729391
Files: 12
Total size: 42.5 KB

Directory structure:
gitextract_osycg9h1/

├── README.md
├── contributors.txt
├── doc/
│   ├── NotebookNavigator.txt
│   └── tags
├── lua/
│   └── notebook-navigator/
│       ├── commenters.lua
│       ├── core.lua
│       ├── highlight.lua
│       ├── init.lua
│       ├── miniai_spec.lua
│       ├── repls.lua
│       └── utils.lua
└── stylua.toml

================================================
FILE CONTENTS
================================================

================================================
FILE: README.md
================================================
# 🚢 Notebook Navigator 🚢

Notebook Navigator lets you manipulate and send code cells to a REPL.

A great feature that comes on by default with VSCode is the ability to define
code cells and send them to a REPL like you would do in a Jupyter notebook but
without the hassle of notebook files. Notebook Navigator brings you back that
functionality and more!

Notebook Navigator comes with the following functions and features:
- Jump up/down between cells
- Run cells (with and without jumping to the next one)
- Create cells above/below the current one
- Split cells
- Comment whole cells
- A [mini.ai](https://github.com/echasnovski/mini.nvim/blob/main/readmes/mini-ai.md) textobject
  specification that you can use standalone
- A [Hydra](https://github.com/anuvyklack/hydra.nvim) mode to quickly manipulate and run
  cells
- Code cell marker highlighting
- Support for multiple languages. Notebooks are not just for Pythonistas!
- ... [and more](#full-api)

This plugin also pairs really well with tools like Jupytext that allow you to
convert easily between `ipynb` and `py` files. For this you may want to use a
plugin such as [jupytext.nvim](https://github.com/GCBallesteros/jupytext.nvim).

![notebook-navigator](assets/notebook_navigator.gif)

This plugin is an evolution of my previous setup which you can find
[here](https://www.maxwellrules.com/misc/nvim_jupyter.html).

## What is a code cell?
A code cell is any code between a cell marker, usually a specially designated comment
and the next cell marker or the end of the buffer. The first line of a buffer has an
implicit cell marker before it.

For example here are a bunch of cells on a Python script
```python
print("Cell 1")
# %%
print("This is cell 2!")
# %%
print("This is the last cell!")
```


## Installation
Here is my [lazy.nvim](https://www.github.com/folke/lazy.nvim) specification for Notebook
Navigator.

I personally like to have the moving between cell commands and cell executing functions
available through leader keymaps but will turn to the Hydra head when many cells need to
be run (just by smashing `x`) or for less commonly used functionality.
```lua
{
  "GCBallesteros/NotebookNavigator.nvim",
  keys = {
    { "]h", function() require("notebook-navigator").move_cell "d" end },
    { "[h", function() require("notebook-navigator").move_cell "u" end },
    { "<leader>X", "<cmd>lua require('notebook-navigator').run_cell()<cr>" },
    { "<leader>x", "<cmd>lua require('notebook-navigator').run_and_move()<cr>" },
  },
  dependencies = {
    "echasnovski/mini.comment",
    "hkupty/iron.nvim", -- repl provider
    -- "akinsho/toggleterm.nvim", -- alternative repl provider
    -- "benlubas/molten-nvim", -- alternative repl provider
    "anuvyklack/hydra.nvim",
  },
  event = "VeryLazy",
  config = function()
    local nn = require "notebook-navigator"
    nn.setup({ activate_hydra_keys = "<leader>h" })
  end,
}
```

## Enabling Mini.hipatterns cell highlighting
The lines delimiting the code cells can have pretty highlighting if you install 
[mini.hipatterns](https://github.com/echasnovski/mini.hipatterns). To activate them you will
have to add an entry into the `highlighters` option of 'mini.hipatterns'. If you are using
[lazy.nvim](https//www.github.com/folke/lazy.nvim) `minihipatterns_spec` and your __only__
configuration was meant to activate cell highlighting then your 'mini.hipatterns' could
look like:

```lua
return {
  "echasnovski/mini.hipatterns",
  event = "VeryLazy",
  dependencies = { "GCBallesteros/NotebookNavigator.nvim" },
  opts = function()
    local nn = require "notebook-navigator"

    local opts = { highlighters = { cells = nn.minihipatterns_spec } }
    return opts
  end,
}
```

If you are after a more simple solution that doesn't require new plugins and looks
more minimal just set the `syntax_highlight` option to `true`.


## Mini.ai integration
The `miniai_spec` function is also a valid mini.ai textobject specification.
All you need to do to add it is to add the `custom_textobjects`  to your 'mini.ai' setup. If you
are using [layz.nvim](https://www.github.com/folke/lazy.nvim) and your __only__ configuration was
meant include the _code cell_ text object then your 'mini.ai' could look like:

```lua
return {
  "echasnovski/mini.ai",
  event = "VeryLazy",
  dependencies = { "GCBallesteros/NotebookNavigator.nvim" },
  opts = function()
    local nn = require "notebook-navigator"

    local opts = { custom_textobjects = { h = nn.miniai_spec } }
    return opts
  end,
}
```


## Detailed configuration
Any options that are not specified when calling `setup` will take on their default values.
```lua
{
  -- Code cell marker. Cells start with the marker and end either at the beginning
  -- of the next cell or at the end of the file.
  -- By default, uses language-specific double percent comments like `# %%`.
  -- This can be overridden for each language with this setting.
  cell_markers = {
    -- python = "# %%",
  },
  -- If not `nil` the keymap defined in the string will activate the hydra head.
  -- If you don't want to use hydra you don't need to install it either.
  activate_hydra_keys = nil,
  -- If `true` a hint panel will be shown when the hydra head is active. If `false`
  -- you get a minimalistic hint on the command line.
  show_hydra_hint = true,
  -- Mappings while the hydra head is active.
  -- Any of the mappings can be set to "nil", the string! Not the value! to unamp it
  hydra_keys = {
    comment = "c",
    run = "X",
    run_and_move = "x",
    move_up = "k",
    move_down = "j",
    add_cell_before = "a",
    add_cell_after = "b",
  },
  -- The repl plugin with which to interface
  -- Current options: "iron" for iron.nvim, "toggleterm" for toggleterm.nvim,
  -- "molten" for molten-nvim or "auto" which checks which of the above are 
  -- installed
  repl_provider = "auto",
  -- Syntax based highlighting. If you don't want to install mini.hipattners or
  -- enjoy a more minimalistic look
  syntax_highlight = false,
  -- (Optional) for use with `mini.hipatterns` to highlight cell markers
  cell_highlight_group = "Folded",
}
```

## Current limitations
If any key gets remapped or unmapped to a different key you will need to set `show_hydra_hint`
to `false`. See issue for more details.


## Dependencies
The currently supported REPLs are:
- [iron.nvim](https://github.com/Vigemus/iron.nvim),
- [toggleterm.nvim](https://github.com/akinsho/toggleterm.nvim) or
- [molten-nvim](https://github.com/benlubas/molten-nvim)

The latter are automatically detected. Support for others like `conjure`
or `yarepl` may be added if people want them or are willing to send in PRs.

Commenting cells of code depends on an external plugin. Either
[comment.nvim](https://github.com/numToStr/Comment.nvim) or
[mini.comment](https://github.com/echasnovski/mini.comment) the two most
popular choices by quite a bit. If you want support for more PRs are welcome.

'mini.ai' is not a dependency but if you want to use the provided
textobject specification (highly recommended) you will then need to have it
installed.

Finally, 'mini.hipatterns' is also not a dependency but can provide line
highlighting to distinguish cell markers from the rest of the text.

## Yanking/Deleting cells
If you setup the mini.ai integration (see below) you can then do things like,
`dah` to delete a cell, `yih` to copy just the code or `vah` to select the full
cell in visual mode. (y)ank, (d)elete and (v)isual also work while inside the
Hydra mode!

## Full API

NotebookNavigator provides more functionality to manipulate cells than it is
directly exposed on the Hydra mode. You may use any of those as additional
keymaps on the plugin configuration or even map them on they Hydra mode as long
as you take heed of the [advice given above](#current-limitations).

- `move_cell(dir)`: Move up or done a cell in the `u`p or `d`own direction.
- `run_cell(repl_args)`: Run the current cell. You may optionally pass a table
of `repl_args` that will be forwarded to the repl. For the details of what is
forwarded exactly and how it is used check `repls.lua` and look for your repl
provider.
- `run_and_move(repl_args)`: Same as above but also move down to the next cell.
- `swap_cell(dir)`: Swap the current cell with the cell above (`dir='u'`) or
below (`dir='d'`).
- `run_all_cells(repl_args)`: Run all the file.
- `run_cells_below(repl_args)`: Run the current cell and all of the ones below.
- `comment_cell`: Comment the code inside the current cell.
- `add_cell_below`: Add a cell marker below the current cell.
- `add_cell_after`: Same as above (deprecated).
- `add_cell_above`: Add a cell marker above the current cell.
- `add_cell_before`: Same as above (deprecated).
- `split_cell`: Add a cell marker at the current line effectively splitting the
cell.
- `merge_cell`: Merge the current cell ith the one above (`dir='u'`) or below
(`dir='d'`)

## Related plugins

- [nvim-various-textobjs](https://github.com/chrisgrieser/nvim-various-textobjs)
also provides a notebook cell object that you might want to consider if you are
not interested on all the other functionality in NotebookNavigator.

## Contributors

A list of contributors can be found on `contributors.txt`.

Special shoutout to `@cnrobertson` who contributed a significant number of
features, some of which are not reflected on the commit history.


================================================
FILE: contributors.txt
================================================
Guillem Ballesteros (@GCBallesteros)
Connor Robertson (@cnrrobertson)
Thomas M Kehrenberg (@tmke8)
Gregory Power (@gregorywaynepower)
pseudometa (@chrisgrieser)
Tyler Baur (@TillerBurr)


================================================
FILE: doc/NotebookNavigator.txt
================================================
==============================================================================
------------------------------------------------------------------------------
*NotebookNavigator.doc* Easily navigate and work with code cells
*NotebookNavigator*

MIT License Copyright (c) 2023 Guillem Ballesteros

==============================================================================

Jupyter notebooks are great for prototyping and quickly iterating on an idea
but the are hard to version control and track how code has been executed. Both
issues are much less problematic when working with scripts. VSCode does this
better. This plugin attempts to bring the functionality of Jupyter and VSCode
style to neovim.

# What is a code cell?
A code cell is any code between a cell marker, usually a specially designated comment
and the next cell marker or the end of the buffer. The first line of a buffer has an
implicit cell marker before it.

# What comes bundled?~
- Jump up/down between cells
- Run cells (with and without jumping to the next one)
- Create cells above/below the current one
- Comment whole cells
- Split cells
- A mini.ai textobject specification that you can use standalone
- A Hydra mode to quickly manipulate and run cells
- Support for multiple languages

# Setup~
Just run `require("notebook-navigator").setup(opts)` as you would with most Lua
nvim packages. Any options that are left unspecified will take on their default
values.

------------------------------------------------------------------------------
                                                               *M.miniai_spec()*
                            `M.miniai_spec`({opts})
Returns the boundaries of the current code cell.

Parameters~
{opts} `(string)` Either "i" to select the inner lines of the cell or "a" for
  the outer cell.

Return~
`(table)` Table with keys from/to indicating the start and end of the cell.
  The from/to fields themselves have a line and col field.

------------------------------------------------------------------------------
                                                                 *M.move_cell()*
                              `M.move_cell`({dir})
Move between cells

Move between cells and indicate wheter we are at the first or last cell via the
string output.

Parameters~
{dir} `(string)` Movement direction. "d" for down and "u" for up.

Return~
`(string)` If movement failed return "first" or "last" if we where at the
  first/last cell.

------------------------------------------------------------------------------
                                                                  *M.run_cell()*
                           `M.run_cell`({repl_args})
Run the current cell under the cursor

Parameters~
{repl_args} `(table|nil)` Optional config for the repl.

------------------------------------------------------------------------------
                                                              *M.run_and_move()*
                         `M.run_and_move`({repl_args})
Run the current cell under the cursor and jump to next cell. If no next cell
is available it will create one like Jupyter notebooks.

Parameters~
{repl_args} `(table|nil)` Optional config for the repl.

------------------------------------------------------------------------------
                                                                 *M.swap_cell()*
                              `M.swap_cell`({dir})
Swap the current cell with the cell immediately above or below

Swap cell with the above or below

Parameters~
{dir} `(string)` Swap direction. "d" for down and "u" for up.

------------------------------------------------------------------------------
                                                                *M.merge_cell()*
                             `M.merge_cell`({dir})
Merge cell

Merge cell with the above or below

Parameters~
{dir} `(string)` Merge direction. "d" for down and "u" for up.

------------------------------------------------------------------------------
                                                             *M.run_all_cells()*
                         `M.run_all_cells`({repl_args})
Run all cells in the file

Parameters~
{repl_args} `(table|nil)` Optional config for the repl.

------------------------------------------------------------------------------
                                                           *M.run_cells_below()*
                        `M.run_cells_below`({repl_args})
Run all cells below (including current cell)

Parameters~
{repl_args} `(table|nil)` Optional config for the repl.

------------------------------------------------------------------------------
                                                              *M.comment_cell()*
                               `M.comment_cell`()
Comment all the contents of the cell under the cursor

The commenting functionality is supported by external plugins. Currently the
following are supported:
- mini.comment
- comment.nvim

------------------------------------------------------------------------------
                                                            *M.add_cell_after()*
                              `M.add_cell_after`()
[Deprecated] Create a cell under the current one and move to it

------------------------------------------------------------------------------
                                                           *M.add_cell_before()*
                             `M.add_cell_before`()
[Deperecated] Create a cell on top of the current one and move to it

------------------------------------------------------------------------------
                                                            *M.add_cell_below()*
                              `M.add_cell_below`()
Create a cell under the current one and move to it

------------------------------------------------------------------------------
                                                            *M.add_cell_above()*
                              `M.add_cell_above`()
Create a cell on top of the current one and move to it

------------------------------------------------------------------------------
                                                                *M.split_cell()*
                                `M.split_cell`()
Spit the cell at the current position by inserting a cell marker

------------------------------------------------------------------------------
                                                                      *M.config*
                                   `M.config`
Module config

Default values:
>
  M.config = {
    -- Code cell marker. Cells start with the marker and end either at the beginning
    -- of the next cell or at the end of the file.
    -- By default, uses language-specific double percent comments like `# %%`.
    -- This can be overridden for each language with this setting.
    cell_markers = {
      -- python = "# %%",
    },

    -- If not `nil` the keymap defined in the string will activate the hydra head
    activate_hydra_keys = nil,
    -- If `true` a hint panel will be shown when the hydra head is active
    show_hydra_hint = true,
    -- Mappings while the hydra head is active.
    hydra_keys = {
      comment = "c",
      run = "X",
      run_and_move = "x",
      move_up = "k",
      move_down = "j",
      add_cell_before = "a",
      add_cell_after = "b",
      split_cell = "s",
    },
    -- The repl plugin with which to interface
    -- Current options: "iron" for iron.nvim, "toggleterm" for toggleterm.nvim,
    -- "molten" for molten-nvim or "auto" which checks which of the above are 
    -- installed
    repl_provider = "auto",
    -- Syntax based highlighting. If you don't want to install mini.hipattners or
    -- enjoy a more minimalistic look
    syntax_highlight = false,
    -- (Optional) for use with `mini.hipatterns` to highlight cell markers
    cell_highlight_group = "Folded",
  }
<

------------------------------------------------------------------------------
                                                                     *M.setup()*
                              `M.setup`({config})
Module setup

Any of the `hydra_keys` mappings can be set to `nil` in order to stop them
from being mapped.
Parameters~
{config} `(table|nil)` Module config table. See |NotebookNavigator.config|.

Usage~
`require('cell-navigator').setup({})` (replace `{}` with your `config` table)
   any config parameter which you not pass will take on its default value.


 vim:tw=78:ts=8:noet:ft=help:norl:

================================================
FILE: doc/tags
================================================
M.add_cell_above()	NotebookNavigator.txt	/*M.add_cell_above()*
M.add_cell_after()	NotebookNavigator.txt	/*M.add_cell_after()*
M.add_cell_before()	NotebookNavigator.txt	/*M.add_cell_before()*
M.add_cell_below()	NotebookNavigator.txt	/*M.add_cell_below()*
M.comment_cell()	NotebookNavigator.txt	/*M.comment_cell()*
M.config	NotebookNavigator.txt	/*M.config*
M.merge_cell()	NotebookNavigator.txt	/*M.merge_cell()*
M.miniai_spec()	NotebookNavigator.txt	/*M.miniai_spec()*
M.move_cell()	NotebookNavigator.txt	/*M.move_cell()*
M.run_all_cells()	NotebookNavigator.txt	/*M.run_all_cells()*
M.run_and_move()	NotebookNavigator.txt	/*M.run_and_move()*
M.run_cell()	NotebookNavigator.txt	/*M.run_cell()*
M.run_cells_below()	NotebookNavigator.txt	/*M.run_cells_below()*
M.setup()	NotebookNavigator.txt	/*M.setup()*
M.split_cell()	NotebookNavigator.txt	/*M.split_cell()*
M.swap_cell()	NotebookNavigator.txt	/*M.swap_cell()*
NotebookNavigator	NotebookNavigator.txt	/*NotebookNavigator*
NotebookNavigator.doc	NotebookNavigator.txt	/*NotebookNavigator.doc*


================================================
FILE: lua/notebook-navigator/commenters.lua
================================================
local commenters = {}

-- comment.nvim
commenters.comment_nvim = function(cell_object)
  local comment = require "Comment.api"
  local curr_pos = vim.api.nvim_win_get_cursor(0)
  local n_lines = cell_object.to.line - cell_object.from.line + 1

  vim.api.nvim_win_set_cursor(0, { cell_object.from.line, 0 })
  comment.toggle.linewise.count(n_lines)
  vim.api.nvim_win_set_cursor(0, curr_pos)
end

-- mini.comment
commenters.mini_comment = function(cell_object)
  local comment = require "mini.comment"
  comment.toggle_lines(cell_object.from.line, cell_object.to.line)
end

-- no recognized comment plugin
commenters.no_comments = function(_)
  vim.notify "[Notebook Navigator] No supported comment plugin available"
end

local has_mini_comment, _ = pcall(require, "mini.comment")
local has_comment_nvim, _ = pcall(require, "Comment.api")
local commenter
if has_mini_comment then
  commenter = commenters["mini_comment"]
elseif has_comment_nvim then
  commenter = commenters["comment_nvim"]
else
  commenter = commenters["no_comments"]
end

return commenter


================================================
FILE: lua/notebook-navigator/core.lua
================================================
local commenter = require "notebook-navigator.commenters"
local get_repl = require "notebook-navigator.repls"
local miniai_spec = require("notebook-navigator.miniai_spec").miniai_spec

local M = {}

M.move_cell = function(dir, cell_marker)
  local search_res
  local result

  if dir == "d" then
    search_res = vim.fn.search("^" .. cell_marker, "W")
    if search_res == 0 then
      result = "last"
    end
  else
    search_res = vim.fn.search("^" .. cell_marker, "bW")
    if search_res == 0 then
      result = "first"
      vim.api.nvim_win_set_cursor(0, { 1, 0 })
    end
  end

  return result
end

M.swap_cell = function(dir, cell_marker)
  local buf_length = vim.api.nvim_buf_line_count(0)
  local should_insert_marker = false

  -- Get cells in their future order
  local starting_cursor = vim.api.nvim_win_get_cursor(0)
  local first_cell
  local second_cell
  if dir == "d" then
    second_cell = miniai_spec("a", cell_marker)
    if second_cell.to.line + 1 > buf_length then
      return
    end
    vim.api.nvim_win_set_cursor(0, { second_cell.to.line + 2, 0 })
    first_cell = miniai_spec("a", cell_marker)
  else
    first_cell = miniai_spec("a", cell_marker)
    if first_cell.from.line - 1 < 1 then
      return
    end
    vim.api.nvim_win_set_cursor(0, { first_cell.from.line - 1, 0 })
    second_cell = miniai_spec("a", cell_marker)

    -- The first cell may not have a marker. If this is the case and we attempt to
    -- swap it down we will be in trouble. In that case we first insert a marker at
    -- the top.
    -- If the line does not start with the cell_marker with set a marker to add
    -- the line later on.
    local first_cell_line =
      vim.api.nvim_buf_get_lines(0, second_cell.from.line - 1, second_cell.from.line, false)[1]

    if string.sub(first_cell_line, 1, string.len(cell_marker)) ~= cell_marker then
      should_insert_marker = true
    end
    --
  end

  -- Combine cells and set in place
  local first_lines = vim.api.nvim_buf_get_lines(0, first_cell.from.line - 1, first_cell.to.line, false)
  local second_lines = vim.api.nvim_buf_get_lines(0, second_cell.from.line - 1, second_cell.to.line, false)

  local final_lines = {}

  for _, v in ipairs(first_lines) do
    table.insert(final_lines, v)
  end

  -- This extra marker protects us agains malformed notebooks that don't have a cell
  -- marker at the top of the file. See the "up" case a few lines above.
  if should_insert_marker then
    table.insert(final_lines, cell_marker)
  end
  for _, v in ipairs(second_lines) do
    table.insert(final_lines, v)
  end
  vim.api.nvim_buf_set_lines(0, second_cell.from.line - 1, first_cell.to.line, false, final_lines)

  -- Put cursor in previous position
  local new_cursor = starting_cursor
  if dir == "d" then
    new_cursor[1] = new_cursor[1] + (first_cell.to.line - first_cell.from.line + 1)
  else
    new_cursor[1] = new_cursor[1] - (second_cell.to.line - second_cell.from.line + 1)
  end
  vim.api.nvim_win_set_cursor(0, new_cursor)
end

M.run_cell = function(cell_marker, repl_provider, repl_args)
  repl_args = repl_args or nil
  repl_provider = repl_provider or "auto"
  local cell_object = miniai_spec("i", cell_marker)

  -- protect ourselves against the case with no actual lines of code
  local n_lines = cell_object.to.line - cell_object.from.line + 1
  if n_lines < 1 then
    return nil
  end

  local repl = get_repl(repl_provider)
  return repl(cell_object.from.line, cell_object.to.line, repl_args, cell_marker)
end

M.run_and_move = function(cell_marker, repl_provider, repl_args)
  local success = M.run_cell(cell_marker, repl_provider, repl_args)

  if success then
    local is_last_cell = M.move_cell("d", cell_marker) == "last"

    -- insert a new cell to replicate the behaviour of jupyter notebooks
    if is_last_cell then
      vim.api.nvim_buf_set_lines(0, -1, -1, false, { cell_marker, "" })
      -- and move to it
      M.move_cell("d", cell_marker)
    end
  end
end

M.run_all_cells = function(repl_provider, repl_args)
  local buf_length = vim.api.nvim_buf_line_count(0)

  local repl = get_repl(repl_provider)
  repl(1, buf_length, repl_args)
end

M.run_cells_below = function(cell_marker, repl_provider, repl_args)
  local buf_length = vim.api.nvim_buf_line_count(0)
  local cell_object = miniai_spec("i", cell_marker)

  local repl = get_repl(repl_provider)
  repl(cell_object.from.line, buf_length, repl_args)
end

M.merge_cell = function(dir, cell_marker)
  local search_res
  local result

  if dir == "d" then
    search_res = vim.fn.search("^" .. cell_marker, "nW")
    vim.api.nvim_buf_set_lines(0, search_res - 1, search_res, false, { "" })
  else
    search_res = vim.fn.search("^" .. cell_marker, "nbW")
    if search_res == 0 then
      return "first"
    else
      vim.api.nvim_buf_set_lines(0, search_res - 1, search_res, false, { "" })
    end
  end

  return result
end

M.comment_cell = function(cell_marker)
  local cell_object = miniai_spec("i", cell_marker)

  -- protect against empty cells
  local n_lines = cell_object.to.line - cell_object.from.line + 1
  if n_lines < 1 then
    return nil
  end
  commenter(cell_object)
end

M.add_cell_below = function(cell_marker)
  local cell_object = miniai_spec("a", cell_marker)

  vim.api.nvim_buf_set_lines(0, cell_object.to.line, cell_object.to.line, false, { cell_marker, "" })
  M.move_cell("d", cell_marker)
end

M.add_cell_above = function(cell_marker)
  local cell_object = miniai_spec("a", cell_marker)

  -- What to do on malformed notebooks? I.e. with no upper cell marker? are they malformed?
  -- What if we have a jupytext header? Code doesn't start at top of buffer.
  vim.api.nvim_buf_set_lines(
    0,
    cell_object.from.line - 1,
    cell_object.from.line - 1,
    false,
    { cell_marker, "" }
  )
  M.move_cell("u", cell_marker)
end

-- We keep this two for backwards compatibility but the prefered way is to use
-- the above/below functions for consistency with jupyter nomenclature
M.add_cell_before = function(cell_marker)
  M.add_cell_above(cell_marker)
end

M.add_cell_after = function(cell_marker)
  M.add_cell_below(cell_marker)
end

M.split_cell = function(cell_marker)
  local cursor_line = vim.api.nvim_win_get_cursor(0)[1]
  vim.api.nvim_buf_set_lines(0, cursor_line - 1, cursor_line - 1, false, { cell_marker })
  vim.api.nvim_win_set_cursor(0, { cursor_line + 1, 0 })
end

return M


================================================
FILE: lua/notebook-navigator/highlight.lua
================================================
local core = require "notebook-navigator.core"
local utils = require "notebook-navigator.utils"

local highlight = {}

highlight.minihipatterns_spec = function(cell_markers, hl_group)
  local notebook_cells = {
    pattern = function(buf_id)
      local cell_marker = utils.get_cell_marker(buf_id, cell_markers)
      if cell_marker then
        local regex_cell_marker = string.gsub("^" .. cell_marker, "%%", "%%%%")
        return regex_cell_marker
      else
        return nil
      end
    end,
    group = "",
    extmark_opts = {
      virt_text = {
        {
          "───────────────────────────────────────────────────────────────",
          hl_group,
        },
      },
      line_hl_group = hl_group,
      hl_eol = true,
    },
  }
  return notebook_cells
end

highlight.setup_autocmd_syntax_highlights = function(cell_markers, hl_group)
  vim.api.nvim_create_augroup("NotebookNavigator", { clear = true })

  -- Create autocmd for every language
  for ft, marker in pairs(cell_markers) do
    local syntax_rule = [[ /^\s*]] .. marker .. [[.*$/]]
    local syntax_cmd = "syntax match CodeCell" .. syntax_rule
    vim.api.nvim_create_autocmd("FileType", {
      pattern = ft,
      group = "NotebookNavigator",
      command = syntax_cmd,
    })
  end
  vim.api.nvim_set_hl(0, "CodeCell", { link = hl_group })
  vim.api.nvim_exec_autocmds("FileType", { group = "NotebookNavigator" })
end

return highlight


================================================
FILE: lua/notebook-navigator/init.lua
================================================
--- *NotebookNavigator.doc* Easily navigate and work with code cells
--- *NotebookNavigator*
---
--- MIT License Copyright (c) 2023 Guillem Ballesteros
---
--- ==============================================================================
---
--- Jupyter notebooks are great for prototyping and quickly iterating on an idea
--- but the are hard to version control and track how code has been executed. Both
--- issues are much less problematic when working with scripts. VSCode does this
--- better. This plugin attempts to bring the functionality of Jupyter and VSCode
--- style to neovim.
---
--- # What is a code cell?
--- A code cell is any code between a cell marker, usually a specially designated comment
--- and the next cell marker or the end of the buffer. The first line of a buffer has an
--- implicit cell marker before it.
---
--- # What comes bundled?~
--- - Jump up/down between cells
--- - Run cells (with and without jumping to the next one)
--- - Create cells above/below the current one
--- - Comment whole cells
--- - Split cells
--- - A mini.ai textobject specification that you can use standalone
--- - A Hydra mode to quickly manipulate and run cells
--- - Support for multiple languages
---
--- # Setup~
--- Just run `require("notebook-navigator").setup(opts)` as you would with most Lua
--- nvim packages. Any options that are left unspecified will take on their default
--- values.

local M = {}

local got_hydra, hydra = pcall(require, "hydra")

local core = require "notebook-navigator.core"
local highlight = require "notebook-navigator.highlight"
local miniai_spec = require("notebook-navigator.miniai_spec").miniai_spec
local utils = require "notebook-navigator.utils"

local cell_marker = function()
  return utils.get_cell_marker(0, M.config.cell_markers)
end

-- Export directly the core functions specialized for cell markers

--- Returns the boundaries of the current code cell.
---
---@param opts string Either "i" to select the inner lines of the cell or "a" for
---   the outer cell.
---
---@return table Table with keys from/to indicating the start and end of the cell.
---   The from/to fields themselves have a line and col field.
M.miniai_spec = function(opts)
  return miniai_spec(opts, cell_marker())
end

--- Move between cells
---
--- Move between cells and indicate wheter we are at the first or last cell via the
--- string output.
---
---@param dir string Movement direction. "d" for down and "u" for up.
---
---@return string If movement failed return "first" or "last" if we where at the
---   first/last cell.
M.move_cell = function(dir)
  return core.move_cell(dir, cell_marker())
end

--- Run the current cell under the cursor
---
---@param repl_args table|nil Optional config for the repl.
M.run_cell = function(repl_args)
  core.run_cell(cell_marker(), M.config.repl_provider, repl_args)
end

--- Run the current cell under the cursor and jump to next cell. If no next cell
--- is available it will create one like Jupyter notebooks.
---
---@param repl_args table|nil Optional config for the repl.
M.run_and_move = function(repl_args)
  core.run_and_move(cell_marker(), M.config.repl_provider, repl_args)
end

--- Swap the current cell with the cell immediately above or below
---
--- Swap cell with the above or below
---
---@param dir string Swap direction. "d" for down and "u" for up.
M.swap_cell = function(dir)
  return core.swap_cell(dir, cell_marker())
end

--- Merge cell
---
--- Merge cell with the above or below
---
---@param dir string Merge direction. "d" for down and "u" for up.
M.merge_cell = function(dir)
  return core.merge_cell(dir, cell_marker())
end

--- Run all cells in the file
---
---@param repl_args table|nil Optional config for the repl.
M.run_all_cells = function(repl_args)
  core.run_all_cells(M.config.repl_provider, repl_args)
end

--- Run all cells below (including current cell)
---
---@param repl_args table|nil Optional config for the repl.
M.run_cells_below = function(repl_args)
  core.run_cells_below(cell_marker(), M.config.repl_provider, repl_args)
end

--- Comment all the contents of the cell under the cursor
---
--- The commenting functionality is supported by external plugins. Currently the
--- following are supported:
--- - mini.comment
--- - comment.nvim
M.comment_cell = function()
  core.comment_cell(cell_marker())
end

--- [Deprecated] Create a cell under the current one and move to it
M.add_cell_after = function()
  core.add_cell_after(cell_marker())
end

--- [Deperecated] Create a cell on top of the current one and move to it
M.add_cell_before = function()
  core.add_cell_before(cell_marker())
end

--- Create a cell under the current one and move to it
M.add_cell_below = function()
  core.add_cell_below(cell_marker())
end

--- Create a cell on top of the current one and move to it
M.add_cell_above = function()
  core.add_cell_above(cell_marker())
end

--- Spit the cell at the current position by inserting a cell marker
M.split_cell = function()
  core.split_cell(cell_marker())
end

local hydra_hint = [[
 _j_/_k_: move down/up   _c_: comment     _a_/_b_: add cell above/below
_x_: run & move down  _s_: split cell   _X_: run
                    _<esc>_/_q_: exit
]]

local function activate_hydra(config)
  -- `hydra_heads` contains all the potential actions that our hydra will be able
  -- to execute. After these definitions the list will get filtered down by checking
  -- if the mapped key is nil
  local hydra_heads = {
    {
      config.hydra_keys.move_up,
      function()
        M.move_cell "u"
      end,
      { desc = "Move up" },
    },
    {
      config.hydra_keys.move_down,
      function()
        M.move_cell "d"
      end,
      { desc = "Move down" },
    },
    {
      config.hydra_keys.comment,
      M.comment_cell,
      { desc = "Comment" },
    },
    {
      config.hydra_keys.run,
      M.run_cell,
      { desc = "Run", nowait = true },
    },
    {
      config.hydra_keys.run_and_move,
      M.run_and_move,
      { desc = "Run & Move", nowait = true },
    },
    {
      config.hydra_keys.add_cell_after,
      M.add_cell_below,
      { desc = "Add cell below", nowait = true },
    },
    {
      config.hydra_keys.add_cell_before,
      M.add_cell_above,
      { desc = "Add cell above", nowait = true },
    },
    {
      config.hydra_keys.split_cell,
      M.split_cell,
      { desc = "Split cell", nowait = true },
    },
    { "q", nil, { exit = true, nowait = true, desc = "exit" } },
    { "<esc>", nil, { exit = true, nowait = true, desc = "exit" } },
  }

  local active_hydra_heads = {}
  for _, h in ipairs(hydra_heads) do
    if h[1] ~= "nil" then
      table.insert(active_hydra_heads, h)
    end
  end

  local hydra_config = {
    name = "NotebookNavigator",
    mode = { "n" },
    config = {
      invoke_on_body = true,
      color = "pink",
      hint = { border = "rounded" },
    },
    body = config.activate_hydra_keys,
    heads = active_hydra_heads,
  }
  if config.show_hydra_hint then
    hydra_config.hint = hydra_hint
  end

  hydra(hydra_config)
end

--- Module config
---
--- Default values:
---@eval return MiniDoc.afterlines_to_code(MiniDoc.current.eval_section)
M.config = {
  -- Code cell marker. Cells start with the marker and end either at the beginning
  -- of the next cell or at the end of the file.
  -- By default, uses language-specific double percent comments like `# %%`.
  -- This can be overridden for each language with this setting.
  cell_markers = {
    -- python = "# %%",
  },

  -- If not `nil` the keymap defined in the string will activate the hydra head
  activate_hydra_keys = nil,
  -- If `true` a hint panel will be shown when the hydra head is active
  show_hydra_hint = true,
  -- Mappings while the hydra head is active.
  hydra_keys = {
    comment = "c",
    run = "X",
    run_and_move = "x",
    move_up = "k",
    move_down = "j",
    add_cell_before = "a",
    add_cell_after = "b",
    split_cell = "s",
  },
  -- The repl plugin with which to interface
  -- Current options: "iron" for iron.nvim, "toggleterm" for toggleterm.nvim,
  -- or "auto" which checks which of the above are installed
  repl_provider = "auto",
  -- Syntax based highlighting. If you don't want to install mini.hipattners or
  -- enjoy a more minimalistic look
  syntax_highlight = false,
  -- (Optional) for use with `mini.hipatterns` to highlight cell markers
  cell_highlight_group = "Folded",
}
--minidoc_afterlines_end

--- Module setup
---
--- Any of the `hydra_keys` mappings can be set to `nil` in order to stop them
--- from being mapped.
---@param config table|nil Module config table. See |NotebookNavigator.config|.
---
---@usage `require('cell-navigator').setup({})` (replace `{}` with your `config` table)
---    any config parameter which you not pass will take on its default value.
M.setup = function(config)
  vim.validate({ config = { config, "table", true } })
  M.config = vim.tbl_deep_extend("force", M.config, config or {})

  vim.validate({
    cell_markers = { M.config.cell_markers, "table" },
    activate_hydra_keys = { M.config.activate_hydra_keys, "string", true },
    show_hydra_hint = { M.config.show_hydra_hint, "boolean" },
    hydra_keys = { M.config.hydra_keys, "table" },
  })

  vim.validate({
    ["config.hydra_keys.comment"] = { M.config.hydra_keys.comment, "string" },
    ["config.hydra_keys.run"] = { M.config.hydra_keys.run, "string" },
    ["config.hydra_keys.run_and_move"] = { M.config.hydra_keys.run_and_move, "string" },
    ["config.hydra_keys.move_up"] = { M.config.hydra_keys.move_up, "string" },
    ["config.hydra_keys.move_down"] = { M.config.hydra_keys.move_down, "string" },
    ["config.hydra_keys.add_cell_before"] = { M.config.hydra_keys.add_cell_before, "string" },
    ["config.hydra_keys.add_cell_after"] = { M.config.hydra_keys.add_cell_after, "string" },
  })

  for ft, marker in pairs(M.config.cell_markers) do
    vim.validate({
      ["config.cell_markers." .. ft] = { marker, "string" },
    })
  end

  if (not got_hydra) and (M.config.activate_hydra_keys ~= nil) then
    vim.notify "[NotebookNavigator] Hydra is not available.\nHydra will not be available."
  end

  if #utils.available_repls == 0 then
    vim.notify "[NotebookNavigator] No supported REPLs available.\nMost functionality will error out."
  elseif
    M.config.repl_provider ~= "auto" and not utils.has_value(utils.available_repls, M.config.repl_provider)
  then
    vim.notify("[NotebookNavigator] The requested repl (" .. M.config.repl_provider .. ") is not available.")
  end

  if (M.config.activate_hydra_keys ~= nil) and got_hydra then
    activate_hydra(M.config)
  end

  --- Highlight spec for mini.hipatterns
  M.minihipatterns_spec = highlight.minihipatterns_spec(M.config.cell_markers, M.config.cell_highlight_group)

  -- Apply syntax highlight rule for cell markers
  if M.config.syntax_highlight and M.config.cell_highlight_group ~= nil then
    highlight.setup_autocmd_syntax_highlights(M.config.cell_markers, M.config.cell_highlight_group)
  end
end

return M


================================================
FILE: lua/notebook-navigator/miniai_spec.lua
================================================
local M = {}

M.miniai_spec = function(opts, cell_marker)
  local start_line = vim.fn.search("^" .. cell_marker, "bcnW")

  -- Just in case the notebook is malformed and doesnt  have a cell marker at the start.
  if start_line == 0 then
    start_line = 1
  else
    if opts == "i" then
      start_line = start_line + 1
    end
  end

  local end_line = vim.fn.search("^" .. cell_marker, "nW") - 1
  if end_line == -1 then
    end_line = vim.fn.line "$"
  end

  local last_col = math.max(vim.fn.getline(end_line):len(), 1)

  local from = { line = start_line, col = 1 }
  local to = { line = end_line, col = last_col }

  return { from = from, to = to }
end

return M


================================================
FILE: lua/notebook-navigator/repls.lua
================================================
local repls = {}

local utils = require "notebook-navigator.utils"

-- iron.nvim
---@diagnostic disable-next-line: unused-local
repls.iron = function(start_line, end_line, repl_args, _cell_marker)
  local lines = vim.api.nvim_buf_get_lines(0, start_line - 1, end_line, false)
  require("iron.core").send(nil, lines)

  return true
end

-- toggleterm
---@diagnostic disable-next-line: unused-local
repls.toggleterm = function(start_line, end_line, repl_args, cell_marker)
  local id = 1
  local trim_spaces = false
  if repl_args then
    id = repl_args.id or 1
    trim_spaces = (repl_args.trim_spaces == nil) or repl_args.trim_spaces
  end
  local current_window = vim.api.nvim_get_current_win()
  local lines = vim.api.nvim_buf_get_lines(0, start_line - 1, end_line, false)

  if not lines or not next(lines) then
    return
  end

  for _, line in ipairs(lines) do
    local l = trim_spaces and line:gsub("^%s+", ""):gsub("%s+$", "") or line
    require("toggleterm").exec(l, id)
  end

  -- Jump back with the cursor where we were at the beginning of the selection
  local cursor_line, cursor_col = unpack(vim.api.nvim_win_get_cursor(0))
  vim.api.nvim_set_current_win(current_window)

  vim.api.nvim_win_set_cursor(current_window, { cursor_line, cursor_col })

  return true
end

-- molten
---@diagnostic disable-next-line: unused-local
repls.molten = function(start_line, end_line, repl_args, cell_marker)
  local line_count = vim.api.nvim_buf_line_count(0)

  if line_count < (end_line + 1) then
    vim.api.nvim_buf_set_lines(0, end_line + 1, end_line + 1, false, { cell_marker, "" })
  end

  local ok, _ = pcall(vim.fn.MoltenEvaluateRange, start_line, end_line + 1)
  if not ok then
    vim.cmd "MoltenInit"
    return false
  end

  return true
end

-- no repl
repls.no_repl = function(_) end

local get_repl = function(repl_provider)
  local available_repls = utils.available_repls
  local chosen_repl = nil
  if repl_provider == "auto" then
    for _, r in ipairs(available_repls) do
      chosen_repl = repls[r]
      break
    end
  else
    chosen_repl = repls[repl_provider]
  end

  -- Check if we actuall got out a supported repl
  if chosen_repl == nil then
    vim.notify("[NotebookNavigator] The provided repl, " .. repl_provider .. ", is not supported.")
    chosen_repl = repls["no_repl"]
  end

  return chosen_repl
end

return get_repl


================================================
FILE: lua/notebook-navigator/utils.lua
================================================
local utils = {}

utils.get_cell_marker = function(bufnr, cell_markers)
  local ft = vim.bo[bufnr].filetype

  if ft == nil or ft == "" then
    print "[NotebookNavigator] utils.lua: Empty filetype"
  end

  local user_opt_cell_marker = cell_markers[ft]
  if user_opt_cell_marker then
    return user_opt_cell_marker
  end

  -- use double percent markers as default for cell markers
  -- DOCS https://jupytext.readthedocs.io/en/latest/formats-scripts.html#the-percent-format
  if not vim.bo.commentstring then
    error("There's no cell marker and no commentstring defined for filetype " .. ft)
  end
  local cstring = string.gsub(vim.bo.commentstring, "^%%", "%%%%")
  local double_percent_cell_marker = cstring:format "%%"
  return double_percent_cell_marker
end

local find_supported_repls = function()
  local supported_repls = {
    { name = "iron", module = "iron" },
    { name = "toggleterm", module = "toggleterm" },
    { name = "molten", module = "molten.health" },
  }

  local available_repls = {}
  for _, repl in pairs(supported_repls) do
    if pcall(require, repl.module) then
      available_repls[#available_repls + 1] = repl.name
    end
  end

  return available_repls
end

utils.available_repls = find_supported_repls()

utils.has_value = function(tab, val)
  for _, value in ipairs(tab) do
    if value == val then
      return true
    end
  end

  return false
end

return utils


================================================
FILE: stylua.toml
================================================
indent_type = "Spaces"
indent_width = 2
column_width = 110
call_parentheses = "NoSingleString"