From 730b62e1da44bac0ad07b1f762c446b906e3ff68 Mon Sep 17 00:00:00 2001
From: David Ibia <david.ibia@boxmarshall.com>
Date: Sun, 10 Mar 2024 12:22:33 +0100
Subject: [PATCH] feat(molten.lua): add autocmds to dynamically adjust
 configuration based on file type and Molten status feat(molten.lua): add
 autocmd to export output chunks to Jupyter notebook on write
 feat(molten.lua): add autocmds to import output chunks from Jupyter notebook
 based on kernel feat(molten.lua): add autocmd to initialize Molten buffer and
 import output chunks

---
 lua/absolute/after/molten.lua | 84 +++++++++++++++++++++++++++++++++++
 1 file changed, 84 insertions(+)
 create mode 100644 lua/absolute/after/molten.lua

diff --git a/lua/absolute/after/molten.lua b/lua/absolute/after/molten.lua
new file mode 100644
index 0000000..e41b580
--- /dev/null
+++ b/lua/absolute/after/molten.lua
@@ -0,0 +1,84 @@
+-- change the configuration when editing a python file
+vim.api.nvim_create_autocmd("BufEnter", {
+	pattern = "*.py",
+	callback = function(e)
+		if string.match(e.file, ".otter.") then
+			return
+		end
+		if require("molten.status").initialized() == "Molten" then -- this is kinda a hack...
+			vim.fn.MoltenUpdateOption("virt_lines_off_by_1", false)
+			vim.fn.MoltenUpdateOption("virt_text_output", false)
+		else
+			vim.g.molten_virt_lines_off_by_1 = false
+			vim.g.molten_virt_text_output = false
+		end
+	end,
+})
+
+-- Undo those config changes when we go back to a markdown or quarto file
+vim.api.nvim_create_autocmd("BufEnter", {
+	pattern = { "*.qmd", "*.md", "*.ipynb" },
+	callback = function(e)
+		if string.match(e.file, ".otter.") then
+			return
+		end
+		if require("molten.status").initialized() == "Molten" then
+			vim.fn.MoltenUpdateOption("virt_lines_off_by_1", true)
+			vim.fn.MoltenUpdateOption("virt_text_output", true)
+		else
+			vim.g.molten_virt_lines_off_by_1 = true
+			vim.g.molten_virt_text_output = true
+		end
+	end,
+})
+
+-- automatically export output chunks to a jupyter notebook on write
+vim.api.nvim_create_autocmd("BufWritePost", {
+	pattern = { "*.ipynb" },
+	callback = function()
+		if require("molten.status").initialized() == "Molten" then
+			vim.cmd("MoltenExportOutput!")
+		end
+	end,
+})
+
+-- automatically import output chunks from a jupyter notebook
+-- tries to find a kernel that matches the kernel in the jupyter notebook
+-- falls back to a kernel that matches the name of the active venv (if any)
+local imb = function(e) -- init molten buffer
+	vim.schedule(function()
+		local kernels = vim.fn.MoltenAvailableKernels()
+		local try_kernel_name = function()
+			local metadata = vim.json.decode(io.open(e.file, "r"):read("a"))["metadata"]
+			return metadata.kernelspec.name
+		end
+		local ok, kernel_name = pcall(try_kernel_name)
+		if not ok or not vim.tbl_contains(kernels, kernel_name) then
+			kernel_name = nil
+			local venv = os.getenv("VIRTUAL_ENV")
+			if venv ~= nil then
+				kernel_name = string.match(venv, "/.+/(.+)")
+			end
+		end
+		if kernel_name ~= nil and vim.tbl_contains(kernels, kernel_name) then
+			vim.cmd(("MoltenInit %s"):format(kernel_name))
+		end
+		vim.cmd("MoltenImportOutput")
+	end)
+end
+
+-- automatically import output chunks from a jupyter notebook
+vim.api.nvim_create_autocmd("BufAdd", {
+	pattern = { "*.ipynb" },
+	callback = imb,
+})
+
+-- we have to do this as well so that we catch files opened like nvim ./hi.ipynb
+vim.api.nvim_create_autocmd("BufEnter", {
+	pattern = { "*.ipynb" },
+	callback = function(e)
+		if vim.api.nvim_get_vvar("vim_did_enter") ~= 1 then
+			imb(e)
+		end
+	end,
+})