A Fork of Magma nvim, with improved image rendering, performance, and user experience
import json
import os
from typing import Any, Dict, List, Optional, Tuple

import pynvim
from molten.images import Canvas, get_canvas_given_provider
from molten.io import MoltenIOError, get_default_save_file, load, save
from molten.moltenbuffer import MoltenBuffer
from molten.options import MoltenOptions
from molten.outputbuffer import OutputBuffer
from molten.runtime import get_available_kernels
from molten.utils import DynamicPosition, MoltenException, Span, notify_error, notify_info, nvimui
from pynvim import Nvim


@pynvim.plugin
class Molten:
    nvim: Nvim
    canvas: Optional[Canvas]
    initialized: bool

    highlight_namespace: int
    extmark_namespace: int

    timer: Optional[int]

    options: MoltenOptions

    # list of nvim buf numbers to the MoltenBuffer object that handles that nvim buf
    buffers: Dict[int, MoltenBuffer]
    # list of kernel names to the MoltenBuffer object that handles that kernel
    molten_buffers: Dict[str, MoltenBuffer]

    def __init__(self, nvim: Nvim):
        self.nvim = nvim
        self.initialized = False

        self.canvas = None
        self.buffers = {}
        self.timer = None
        self.molten_buffers = {}

    def _initialize(self) -> None:
        assert not self.initialized

        self.options = MoltenOptions(self.nvim)

        self.canvas = get_canvas_given_provider(self.options.image_provider, self.nvim)
        self.canvas.init()

        self.highlight_namespace = self.nvim.funcs.nvim_create_namespace("molten-highlights")
        self.extmark_namespace = self.nvim.funcs.nvim_create_namespace("molten-extmarks")

        self.timer = self.nvim.eval("timer_start(500, 'MoltenTick', {'repeat': -1})")  # type: ignore

        self._setup_highlights()
        self._set_autocommands()

        self.initialized = True

    def _set_autocommands(self) -> None:
        self.nvim.command("augroup molten")
        self.nvim.command("autocmd CursorMoved  * call MoltenOnCursorMoved()")
        self.nvim.command("autocmd CursorMovedI * call MoltenOnCursorMoved()")
        self.nvim.command("autocmd WinScrolled  * call MoltenOnWinScrolled()")
        self.nvim.command("autocmd BufEnter     * call MoltenUpdateInterface()")
        self.nvim.command("autocmd BufLeave     * call MoltenClearInterface()")
        self.nvim.command("autocmd BufUnload    * call MoltenOnBufferUnload()")
        self.nvim.command("autocmd ExitPre      * call MoltenOnExitPre()")
        self.nvim.command("augroup END")

    def _setup_highlights(self) -> None:
        self.nvim.exec_lua("_hl_utils = require('hl_utils')")
        hl_utils = self.nvim.lua._hl_utils
        hl_utils.set_default_highlights(self.options.hl.defaults)

    def _deinitialize(self) -> None:
        for molten in self.buffers.values():
            molten.deinit()
        if self.canvas is not None:
            self.canvas.deinit()
        if self.timer is not None:
            self.nvim.funcs.timer_stop(self.timer)

    def _initialize_if_necessary(self) -> None:
        if not self.initialized:
            self._initialize()

    def _get_molten(self, requires_instance: bool) -> Optional[MoltenBuffer]:
        maybe_molten = self.buffers.get(self.nvim.current.buffer.number)
        if requires_instance and maybe_molten is None:
            raise MoltenException("Molten is not initialized; run `:MoltenInit` to initialize.")
        return maybe_molten

    def _clear_interface(self) -> None:
        if not self.initialized:
            return

        for moltenbuf in self.buffers.values():
            moltenbuf.clear_interface()
            moltenbuf.clear_open_output_windows()
        assert self.canvas is not None
        self.canvas.present()

    def _update_interface(self) -> None:
        if not self.initialized:
            return

        molten = self._get_molten(False)
        if molten is None:
            return

        molten.update_interface()

    def _on_cursor_moved(self, scrolled=False) -> None:
        if not self.initialized:
            return

        molten = self._get_molten(False)
        if molten is None:
            return

        molten.on_cursor_moved(scrolled)

    def _ask_for_choice(self, preface: str, options: List[str]) -> Optional[str]:
        index = self.nvim.funcs.inputlist(
            [preface] + [f"{i+1}. {option}" for i, option in enumerate(options)]
        )
        if index == 0:
            return None
        else:
            return options[index - 1]

    def _initialize_buffer(self, kernel_name: str, shared=False) -> MoltenBuffer:
        assert self.canvas is not None
        if shared:  # use an existing molten buffer, for a new neovim buffer
            molten = self.molten_buffers.get(kernel_name)
            if molten is not None:
                molten.add_nvim_buffer(self.nvim.current.buffer)
                self.buffers[self.nvim.current.buffer.number] = molten
                return molten

            notify_info(
                self.nvim,
                f"No running kernel {kernel_name} to share. Continuing with a new kernel.",
            )

        molten = MoltenBuffer(
            self.nvim,
            self.canvas,
            self.highlight_namespace,
            self.extmark_namespace,
            self.nvim.current.buffer,
            self.options,
            kernel_name,
        )

        self.molten_buffers[kernel_name] = molten
        self.buffers[self.nvim.current.buffer.number] = molten
        molten._doautocmd("MoltenInitPost")

        return molten

    @pynvim.command("MoltenInit", nargs="*", sync=True, complete="file")  # type: ignore
    @nvimui  # type: ignore
    def command_init(self, args: List[str]) -> None:
        self._initialize_if_necessary()

        shared = False
        if args and args[0] == "shared":
            shared = True
            args = args[1:]

        if args:
            kernel_name = args[0]
            self._initialize_buffer(kernel_name, shared=shared)
        else:
            PROMPT = "Select the kernel to launch:"
            available_kernels = get_available_kernels()

            if len(available_kernels) == 0:
                notify_error(self.nvim, "Unable to find any kernels to launch.")
                return

            if shared:
                # Only show kernels that are already tracked by Molten
                available_kernels = list(
                    filter(lambda x: x in self.molten_buffers.keys(), available_kernels)
                )

            if len(available_kernels) == 0:
                notify_error(
                    self.nvim,
                    "No running kernels to share. Please use :MoltenInit without the shared option.",
                )
                return

            lua = f"""
                vim.schedule_wrap(function()
                    vim.ui.select(
                        {{{", ".join(repr(x) for x in available_kernels)}}},
                        {{prompt = "{PROMPT}"}},
                        function(choice)
                            if choice ~= nil then
                                vim.schedule_wrap(function()
                                    vim.cmd("MoltenInit {'shared ' if shared else ''}" .. choice)
                                end)()
                            end
                        end
                    )
                end)()
            """
            self.nvim.exec_lua(lua, async_=False)

    def _deinit_buffer(self, molten: MoltenBuffer) -> None:
        molten.deinit()
        for buf in molten.buffers:
            del self.buffers[buf.number]

    @pynvim.command("MoltenDeinit", nargs=0, sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_deinit(self) -> None:
        self._initialize_if_necessary()

        molten = self._get_molten(True)
        assert molten is not None

        self._clear_interface()

        self._deinit_buffer(molten)

    def _do_evaluate(self, pos: Tuple[Tuple[int, int], Tuple[int, int]]) -> None:
        self._initialize_if_necessary()

        molten = self._get_molten(True)
        assert molten is not None

        bufno = self.nvim.current.buffer.number
        span = Span(
            DynamicPosition(self.nvim, self.extmark_namespace, bufno, *pos[0]),
            DynamicPosition(self.nvim, self.extmark_namespace, bufno, *pos[1]),
        )

        code = span.get_text(self.nvim)

        molten.run_code(code, span)

    def _do_evaluate_expr(self, expr):
        self._initialize_if_necessary()

        molten = self._get_molten(True)
        assert molten is not None
        bufno = self.nvim.current.buffer.number
        span = Span(
            DynamicPosition(self.nvim, self.extmark_namespace, bufno, 0, 0),
            DynamicPosition(self.nvim, self.extmark_namespace, bufno, 0, 0),
        )
        molten.run_code(expr, span)

    @pynvim.function("MoltenUpdateOption", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def function_update_option(self, args) -> None:
        molten = self._get_molten(True)
        assert molten is not None

        if len(args) == 2:
            option, value = args
            molten.options.update_option(option, value)
        else:
            notify_error(
                self.nvim,
                f"Wrong number of arguments passed to :MoltenUpdateOption, expected 2, given {len(args)}",
            )

    @pynvim.command("MoltenEnterOutput", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_enter_output_window(self) -> None:
        molten = self._get_molten(True)
        assert molten is not None
        molten.enter_output()

    @pynvim.command("MoltenEvaluateArgument", nargs=1, sync=True)
    @nvimui
    def commnand_molten_evaluate_argument(self, expr) -> None:
        assert len(expr) == 1
        self._do_evaluate_expr(expr[0])

    @pynvim.command("MoltenEvaluateVisual", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_evaluate_visual(self) -> None:
        _, lineno_begin, colno_begin, _ = self.nvim.funcs.getpos("'<")
        _, lineno_end, colno_end, _ = self.nvim.funcs.getpos("'>")
        span = (
            (
                lineno_begin - 1,
                min(colno_begin, len(self.nvim.funcs.getline(lineno_begin))) - 1,
            ),
            (
                lineno_end - 1,
                min(colno_end, len(self.nvim.funcs.getline(lineno_end))),
            ),
        )

        self._do_evaluate(span)

    @pynvim.function("MoltenEvaluateRange", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def evaulate_range(self, *args) -> None:
        start_line, end_line = args[0]
        span = (
            (start_line - 1, 0),
            (end_line - 1, len(self.nvim.funcs.getline(end_line))),
        )
        self._do_evaluate(span)

    @pynvim.command("MoltenEvaluateOperator", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_evaluate_operator(self) -> None:
        self._initialize_if_necessary()

        self.nvim.options["operatorfunc"] = "MoltenOperatorfunc"
        self.nvim.feedkeys("g@")

    @pynvim.command("MoltenEvaluateLine", nargs=0, sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_evaluate_line(self) -> None:
        _, lineno, _, _, _ = self.nvim.funcs.getcurpos()
        lineno -= 1

        span = ((lineno, 0), (lineno, -1))

        self._do_evaluate(span)

    @pynvim.command("MoltenReevaluateCell", nargs=0, sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_evaluate_cell(self) -> None:
        self._initialize_if_necessary()

        molten = self._get_molten(True)
        assert molten is not None

        molten.reevaluate_cell()

    @pynvim.command("MoltenInterrupt", nargs=0, sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_interrupt(self) -> None:
        molten = self._get_molten(True)
        assert molten is not None

        molten.interrupt()

    @pynvim.command("MoltenRestart", nargs=0, sync=True, bang=True)  # type: ignore # noqa
    @nvimui  # type: ignore
    def command_restart(self, bang: bool) -> None:
        molten = self._get_molten(True)
        assert molten is not None

        molten.restart(delete_outputs=bang)

    @pynvim.command("MoltenDelete", nargs=0, sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_delete(self) -> None:
        self._initialize_if_necessary()

        molten = self._get_molten(True)
        assert molten is not None

        molten.delete_cell()

    @pynvim.command("MoltenShowOutput", nargs=0, sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_show_output(self) -> None:
        self._initialize_if_necessary()

        molten = self._get_molten(True)
        assert molten is not None

        molten.should_show_display_window = True
        self._update_interface()

    @pynvim.command("MoltenHideOutput", nargs=0, sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_hide_output(self) -> None:
        molten = self._get_molten(False)
        if molten is None:
            # get the current buffer, and then search for it in all molten buffers
            cur_buf = self.nvim.current.buffer
            for moltenbuf in self.buffers.values():
                # if we find it, then we know this is a molten output, and we can safely quit and
                # call hide to hide it
                if cur_buf in map(lambda x: x.display_buf, moltenbuf.outputs.values()):
                    self.nvim.command("q")
                    self.nvim.command(":MoltenHideOutput")
                    return
            return

        molten.should_show_display_window = False
        self._update_interface()

    @pynvim.command("MoltenSave", nargs="?", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_save(self, args: List[str]) -> None:
        self._initialize_if_necessary()

        buf = self.nvim.current.buffer
        if args:
            path = args[0]
        else:
            path = get_default_save_file(self.options, buf)

        dirname = os.path.dirname(path)
        if not os.path.exists(dirname):
            os.makedirs(dirname)

        molten = self._get_molten(True)
        assert molten is not None

        with open(path, "w") as file:
            json.dump(save(molten, buf.number), file)

    @pynvim.command("MoltenLoad", nargs="*", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def command_load(self, args: List[str]) -> None:
        self._initialize_if_necessary()

        shared = False
        if args and args[0] == "shared":
            shared = True
            args = args[1:]

        if len(args) == 1:
            path = args[0]
        else:
            path = get_default_save_file(self.options, self.nvim.current.buffer)

        if self.nvim.current.buffer.number in self.buffers:
            raise MoltenException(
                "Molten is already initialized for this buffer; MoltenLoad initializes Molten."
            )

        with open(path) as file:
            data = json.load(file)

        molten = None

        try:
            MoltenIOError.assert_has_key(data, "version", int)
            if (version := data["version"]) != 1:
                raise MoltenIOError(f"Bad version: {version}")

            MoltenIOError.assert_has_key(data, "kernel", str)
            kernel_name = data["kernel"]

            molten = self._initialize_buffer(kernel_name, shared=shared)

            load(self.nvim, molten, self.nvim.current.buffer, data)

            self._update_interface()
        except MoltenIOError as err:
            if molten is not None:
                self._deinit_buffer(molten)

            raise MoltenException("Error while doing Molten IO: " + str(err))

    # Internal functions which are exposed to VimScript

    @pynvim.function("MoltenClearInterface", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def function_clear_interface(self, _: List[Any]) -> None:
        self._clear_interface()

    @pynvim.function("MoltenOnBufferUnload", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def function_on_buffer_unload(self, _: Any) -> None:
        abuf_str = self.nvim.funcs.expand("<abuf>")
        if not abuf_str:
            return

        molten = self.buffers.get(int(abuf_str))
        if molten is None:
            return

        self._deinit_buffer(molten)

    @pynvim.function("MoltenOnExitPre", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def function_on_exit_pre(self, _: Any) -> None:
        self._deinitialize()

    @pynvim.function("MoltenTick", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def function_molten_tick(self, _: Any) -> None:
        self._initialize_if_necessary()

        molten = self._get_molten(False)
        if molten is None:
            return

        molten.tick()

    @pynvim.function("MoltenUpdateInterface", sync=True)  # type: ignore
    @nvimui  # type: ignore
    def function_update_interface(self, _: Any) -> None:
        self._update_interface()

    @pynvim.function("MoltenOnCursorMoved", sync=True)
    @nvimui
    def function_on_cursor_moved(self, _) -> None:
        self._on_cursor_moved()

    @pynvim.function("MoltenOnWinScrolled", sync=True)
    @nvimui
    def function_on_win_scrolled(self, _) -> None:
        self._on_cursor_moved(scrolled=True)

    @pynvim.function("MoltenOperatorfunc", sync=True)
    @nvimui
    def function_molten_operatorfunc(self, args) -> None:
        if not args:
            return

        kind = args[0]

        _, lineno_begin, colno_begin, _ = self.nvim.funcs.getpos("'[")
        _, lineno_end, colno_end, _ = self.nvim.funcs.getpos("']")

        if kind == "line":
            colno_begin = 1
            colno_end = -1
        elif kind == "char":
            pass
        else:
            raise MoltenException(f"this kind of selection is not supported: '{kind}'")

        span = (
            (
                lineno_begin - 1,
                min(colno_begin, len(self.nvim.funcs.getline(lineno_begin))) - 1,
            ),
            (
                lineno_end - 1,
                min(colno_end, len(self.nvim.funcs.getline(lineno_end))),
            ),
        )

        self._do_evaluate(span)

    @pynvim.function("MoltenDefineCell", sync=True)
    def function_molten_define_cell(self, args: List[int]) -> None:
        if not args:
            return

        self._initialize_if_necessary()
        molten = self._get_molten(True)
        assert molten is not None
        assert self.canvas is not None

        start = args[0]
        end = args[1]
        bufno = self.nvim.current.buffer.number
        span = Span(
            DynamicPosition(self.nvim, self.extmark_namespace, bufno, start - 1, 0),
            DynamicPosition(self.nvim, self.extmark_namespace, bufno, end - 1, -1),
        )
        molten.outputs[span] = OutputBuffer(self.nvim, self.canvas, self.options)

rplugin/python3/molten/__init__.py
benlubas/molten-nvim
from typing import Dict, Set
from abc import ABC, abstractmethod

from pynvim import Nvim, logging


class Canvas(ABC):
    @abstractmethod
    def init(self) -> None:
        """
        Initialize the canvas.

        This will be called before the canvas is ever used.
        """

    @abstractmethod
    def deinit(self) -> None:
        """
        Deinitialize the canvas.

        The canvas will not be used after this operation.
        """

    @abstractmethod
    def present(self) -> None:
        """
        Present the canvas.

        This is called only when a redraw is necessary -- so, if desired, it
        can be implemented so that `clear` and `add_image` only queue images as
        to be drawn, and `present` actually performs the operations, in order
        to reduce flickering.
        """

    @abstractmethod
    def clear(self) -> None:
        """
        Clear all images from the canvas.
        """

    @abstractmethod
    def img_size(self, identifier: str) -> Dict[str, int]:
        """
        Get the height of an image in terminal rows.
        """

    @abstractmethod
    def add_image(
        self,
        path: str,
        x: int,
        y: int,
        bufnr: int,
    ) -> str:
        """
        Add an image to the canvas.

        Parameters
        - path: str
          Path to the image we want to show
        - x: int
          Column number of where the image is supposed to be drawn at (top-left
          corner).
        - y: int
          Row number of where the image is supposed to be drawn at (top-right
          corner).
        - bufnr: int
          The buffer number for the buffer in which to draw the image.

        Returns:
        str the identifier for the image
        """


class NoCanvas(Canvas):
    def __init__(self) -> None:
        pass

    def init(self) -> None:
        pass

    def deinit(self) -> None:
        pass

    def present(self) -> None:
        pass

    def clear(self) -> None:
        pass

    def img_size(self, _indentifier: str) -> Dict[str, int]:
        return {"height": 0, "width": 0}

    def add_image(
        self,
        _path: str,
        _x: int,
        _y: int,
        _window: int,
    ) -> None:
        pass


# I think this class will end up being calls to equivalent lua functions in some lua file
# somewhere
class ImageNvimCanvas(Canvas):
    nvim: Nvim
    to_make_visible: Set[str]
    to_make_invisible: Set[str]
    visible: Set[str]

    def __init__(self, nvim: Nvim):
        self.nvim = nvim
        self.images = {}
        self.visible = set()
        self.to_make_visible = set()
        self.to_make_invisible = set()
        self.next_id = 0

    def init(self) -> None:
        self.nvim.exec_lua("_image = require('load_image_nvim').image_api")
        self.nvim.exec_lua("_image_utils = require('load_image_nvim').image_utils")
        self.image_api = self.nvim.lua._image
        self.image_utils = self.nvim.lua._image_utils

    def deinit(self) -> None:
        self.image_api.clear_all()
        self.images.clear()

    def present(self) -> None:
        # images to both show and hide should be ignored
        to_work_on = self.to_make_visible.difference(
            self.to_make_visible.intersection(self.to_make_invisible)
        )
        self.to_make_invisible.difference_update(self.to_make_visible)
        for identifier in self.to_make_invisible:
            self.image_api.clear(identifier)

        for identifier in to_work_on:
            size = self.img_size(identifier)
            self.image_api.render(identifier, size)

        self.visible.update(self.to_make_visible)
        self.to_make_invisible.clear()
        self.to_make_visible.clear()

    def clear(self) -> None:
        for img in self.visible:
            self.image_api.clear(img)

    def img_size(self, identifier: str) -> Dict[str, int]:
        return self.image_api.image_size(identifier)

    def add_image(
        self,
        path: str,
        x: int,
        y: int,
        bufnr: int,
    ) -> str:
        if path not in self.images:
            img = self.image_api.from_file(
                path,
                {
                    "id": path,
                    "buffer": bufnr,
                    "with_virtual_padding": True,
                    "x": x,
                    "y": y,
                },
            )
            self.to_make_visible.add(img)
            return img
        return path


def get_canvas_given_provider(name: str, nvim: Nvim) -> Canvas:
    if name == "none":
        return NoCanvas()
    elif name == "image.nvim":
        return ImageNvimCanvas(nvim)
    else:
        nvim.api.notify(
            f"[Molten] unknown image provider: `{name}`",
            logging.ERROR,
            {"title": "Molten"},
        )
        return NoCanvas()

rplugin/python3/molten/images.py
benlubas/molten-nvim
from typing import Type, Optional, Dict, Any
import os
from pynvim import Nvim

from pynvim.api import Buffer

from molten.utils import MoltenException, Span, DynamicPosition
from molten.options import MoltenOptions
from molten.outputchunks import OutputStatus, Output, to_outputchunk
from molten.outputbuffer import OutputBuffer
from molten.moltenbuffer import MoltenBuffer


class MoltenIOError(Exception):
    @classmethod
    def assert_has_key(
        cls, data: Dict[str, Any], key: str, type_: Optional[Type[Any]] = None
    ) -> Any:
        if key not in data:
            raise cls(f"Missing key: {key}")
        value = data[key]
        if type_ is not None and not isinstance(value, type_):
            raise cls(
                f"Incorrect type for key '{key}': expected {type_.__name__}, \
                got {type(value).__name__}"
            )
        return value


def get_default_save_file(options: MoltenOptions, buffer: Buffer) -> str:
    # XXX: this is string containment checking. Beware.
    if "nofile" in buffer.options["buftype"]:
        raise MoltenException("Buffer does not correspond to a file")

    mangled_name = buffer.name.replace("%", "%%").replace("/", "%")

    return os.path.join(options.save_path, mangled_name + ".json")


def load(nvim: Nvim, moltenbuffer: MoltenBuffer, nvim_buffer: Buffer, data: Dict[str, Any]) -> None:
    MoltenIOError.assert_has_key(data, "content_checksum", str)

    if moltenbuffer._get_content_checksum() != data["content_checksum"]:
        raise MoltenIOError("Buffer contents' checksum does not match!")

    MoltenIOError.assert_has_key(data, "cells", list)
    for cell in data["cells"]:
        MoltenIOError.assert_has_key(cell, "span", dict)
        MoltenIOError.assert_has_key(cell["span"], "begin", dict)
        MoltenIOError.assert_has_key(cell["span"]["begin"], "lineno", int)
        MoltenIOError.assert_has_key(cell["span"]["begin"], "colno", int)
        MoltenIOError.assert_has_key(cell["span"], "end", dict)
        MoltenIOError.assert_has_key(cell["span"]["end"], "lineno", int)
        MoltenIOError.assert_has_key(cell["span"]["end"], "colno", int)
        begin_position = DynamicPosition(
            moltenbuffer.nvim,
            moltenbuffer.extmark_namespace,
            nvim_buffer.number,
            cell["span"]["begin"]["lineno"],
            cell["span"]["begin"]["colno"],
        )
        end_position = DynamicPosition(
            moltenbuffer.nvim,
            moltenbuffer.extmark_namespace,
            nvim_buffer.number,
            cell["span"]["end"]["lineno"],
            cell["span"]["end"]["colno"],
        )
        span = Span(begin_position, end_position)

        # XXX: do we really want to have the execution count here?
        #      what happens when the counts start to overlap?
        MoltenIOError.assert_has_key(cell, "execution_count", int)
        output = Output(cell["execution_count"])

        MoltenIOError.assert_has_key(cell, "status", int)
        output.status = OutputStatus(cell["status"])

        MoltenIOError.assert_has_key(cell, "success", bool)
        output.success = cell["success"]

        MoltenIOError.assert_has_key(cell, "chunks", list)
        for chunk in cell["chunks"]:
            MoltenIOError.assert_has_key(chunk, "data", dict)
            MoltenIOError.assert_has_key(chunk, "metadata", dict)
            output.chunks.append(
                to_outputchunk(
                    nvim,
                    moltenbuffer.runtime._alloc_file,
                    chunk["data"],
                    chunk["metadata"],
                )
            )

        output.old = True

        moltenbuffer.outputs[span] = OutputBuffer(
            moltenbuffer.nvim, moltenbuffer.canvas, moltenbuffer.options
        )
        moltenbuffer.outputs[span].output = output


def save(moltenbuffer: MoltenBuffer, nvim_buffer: int) -> Dict[str, Any]:
    """ Save the current kernel state for the given buffer. """
    return {
        "version": 1,
        "kernel": moltenbuffer.runtime.kernel_name,
        "content_checksum": moltenbuffer._get_content_checksum(),
        "cells": [
            {
                "span": {
                    "begin": {
                        "lineno": span.begin.lineno,
                        "colno": span.begin.colno,
                    },
                    "end": {
                        "lineno": span.end.lineno,
                        "colno": span.end.colno,
                    },
                },
                "execution_count": output.output.execution_count,
                "status": output.output.status.value,
                "success": output.output.success,
                "chunks": [
                    {
                        "data": chunk.jupyter_data,
                        "metadata": chunk.jupyter_metadata,
                    }
                    for chunk in output.output.chunks
                    if chunk.jupyter_data is not None
                    and chunk.jupyter_metadata is not None
                ],
            }
            for span, output in moltenbuffer.outputs.items() if span.begin.bufno == nvim_buffer
        ],
    }

rplugin/python3/molten/io.py
benlubas/molten-nvim
from typing import List, Optional, Dict
from queue import Queue
import hashlib

from pynvim import Nvim
from pynvim.api import Buffer

from molten.options import MoltenOptions
from molten.images import Canvas
from molten.utils import MoltenException, Position, Span, notify_info
from molten.outputbuffer import OutputBuffer
from molten.outputchunks import OutputStatus
from molten.runtime import JupyterRuntime


class MoltenBuffer:
    nvim: Nvim
    canvas: Canvas
    highlight_namespace: int
    extmark_namespace: int
    buffers: List[Buffer]

    runtime: JupyterRuntime

    outputs: Dict[Span, OutputBuffer]
    current_output: Optional[Span]
    queued_outputs: "Queue[Span]"

    selected_cell: Optional[Span]
    should_show_display_window: bool
    updating_interface: bool

    options: MoltenOptions

    def __init__(
        self,
        nvim: Nvim,
        canvas: Canvas,
        highlight_namespace: int,
        extmark_namespace: int,
        main_buffer: Buffer,
        options: MoltenOptions,
        kernel_name: str,
    ):
        self.nvim = nvim
        self.canvas = canvas
        self.highlight_namespace = highlight_namespace
        self.extmark_namespace = extmark_namespace
        self.buffers = [main_buffer]

        self._doautocmd("MoltenInitPre")

        self.runtime = JupyterRuntime(nvim, kernel_name, options)

        self.outputs = {}
        self.current_output = None
        self.queued_outputs = Queue()

        self.selected_cell = None
        self.should_show_display_window = False
        self.updating_interface = False

        self.options = options

    def _doautocmd(self, autocmd: str) -> None:
        assert " " not in autocmd
        self.nvim.command(f"doautocmd User {autocmd}")

    def add_nvim_buffer(self, buffer: Buffer) -> None:
        self.buffers.append(buffer)

    def deinit(self) -> None:
        self._doautocmd("MoltenDeinitPre")
        self.runtime.deinit()
        self._doautocmd("MoltenDeinitPost")

    def interrupt(self) -> None:
        self.runtime.interrupt()

    def restart(self, delete_outputs: bool = False) -> None:
        if delete_outputs:
            self.outputs = {}
            self.clear_interface()

        self.runtime.restart()

    def run_code(self, code: str, span: Span) -> None:
        self._delete_all_cells_in_span(span)
        self.runtime.run_code(code)

        if span in self.outputs:
            self.outputs[span].clear_interface()
            del self.outputs[span]

        self.outputs[span] = OutputBuffer(self.nvim, self.canvas, self.options)
        self.queued_outputs.put(span)

        self.selected_cell = span
        self.should_show_display_window = True
        self.update_interface()

        self._check_if_done_running()

    def reevaluate_cell(self) -> None:
        self.selected_cell = self._get_selected_span()
        if self.selected_cell is None:
            raise MoltenException("Not in a cell")

        code = self.selected_cell.get_text(self.nvim)

        self.run_code(code, self.selected_cell)

    def _check_if_done_running(self) -> None:
        # TODO: refactor
        is_idle = (self.current_output is None or not self.current_output in self.outputs) or (
            self.current_output is not None
            and self.outputs[self.current_output].output.status == OutputStatus.DONE
        )
        if is_idle and not self.queued_outputs.empty():
            key = self.queued_outputs.get_nowait()
            self.current_output = key

    def tick(self) -> None:
        self._check_if_done_running()

        was_ready = self.runtime.is_ready()
        if self.current_output is None or not self.current_output in self.outputs:
            did_stuff = self.runtime.tick(None)
        else:
            did_stuff = self.runtime.tick(self.outputs[self.current_output].output)
        if did_stuff:
            self.update_interface()
        if not was_ready and self.runtime.is_ready():
            notify_info(self.nvim, f"Kernel '{self.runtime.kernel_name}' is ready.")

    def enter_output(self) -> None:
        if self.selected_cell is not None:
            if self.options.enter_output_behavior != "no_open":
                self.should_show_display_window = True
            self.should_show_display_window = self.outputs[self.selected_cell].enter(self.selected_cell.end)

    def _get_cursor_position(self) -> Position:
        _, lineno, colno, _, _ = self.nvim.funcs.getcurpos()
        return Position(self.nvim.current.buffer.number, lineno - 1, colno - 1)

    def clear_interface(self) -> None:
        if self.updating_interface:
            return

        for buffer in self.buffers:
            self.nvim.funcs.nvim_buf_clear_namespace(
                buffer.number,
                self.highlight_namespace,
                0,
                -1,
            )

    def clear_open_output_windows(self) -> None:
        for output in self.outputs.values():
            output.clear_interface()

    def _get_selected_span(self) -> Optional[Span]:
        current_position = self._get_cursor_position()
        selected = None
        for span in reversed(self.outputs.keys()):
            if current_position in span:
                selected = span
                break

        return selected

    def _delete_all_cells_in_span(self, span: Span) -> None:
        for output_span in reversed(list(self.outputs.keys())):
            if (
                output_span.begin in span
                or output_span.end in span
                or span.begin in output_span
                or span.end in output_span
            ):
                self.outputs[output_span].clear_interface()
                del self.outputs[output_span]

    def delete_cell(self) -> None:
        self.selected_cell = self._get_selected_span()
        if self.selected_cell is None:
            return

        self.outputs[self.selected_cell].clear_interface()
        del self.outputs[self.selected_cell]

    def update_interface(self) -> None:
        buffer_numbers = [buf.number for buf in self.buffers]
        if self.nvim.current.buffer.number not in buffer_numbers:
            return

        # TODO: I think this is redundant
        if self.nvim.current.window.buffer.number not in buffer_numbers:
            return

        self.clear_interface()

        self.updating_interface = True

        selected_cell = self._get_selected_span()

        # Clear the cell we just left
        if self.selected_cell != selected_cell and self.selected_cell is not None:
            self.outputs[self.selected_cell].clear_interface()

        if selected_cell is None:
            self.should_show_display_window = False

        self.selected_cell = selected_cell

        if self.selected_cell is not None:
            self._show_selected(self.selected_cell)
        self.canvas.present()

        self.updating_interface = False

    def on_cursor_moved(self, scrolled=False) -> None:
        selected_cell = self._get_selected_span()

        if (
            self.selected_cell is None
            and selected_cell is not None
            and self.options.auto_open_output
        ):
            self.should_show_display_window = True

        if self.selected_cell == selected_cell and selected_cell is not None:
            if (
                scrolled
                and selected_cell.end.lineno < self.nvim.funcs.line("w$")
                and self.should_show_display_window
            ):
                self.update_interface()
            return

        self.update_interface()

    def _show_selected(self, span: Span) -> None:
        """Show the selected cell. Can only have a selected cell in the current buffer"""
        buf = self.nvim.current.buffer
        if buf.number not in [b.number for b in self.buffers]:
            return

        if span.begin.lineno == span.end.lineno:
            self.nvim.funcs.nvim_buf_add_highlight(
                buf.number,
                self.highlight_namespace,
                self.options.hl.cell,
                span.begin.lineno,
                span.begin.colno,
                span.end.colno,
            )
        else:
            self.nvim.funcs.nvim_buf_add_highlight(
                buf.number,
                self.highlight_namespace,
                self.options.hl.cell,
                span.begin.lineno,
                span.begin.colno,
                -1,
            )
            for lineno in range(span.begin.lineno + 1, span.end.lineno):
                self.nvim.funcs.nvim_buf_add_highlight(
                    buf.number,
                    self.highlight_namespace,
                    self.options.hl.cell,
                    lineno,
                    0,
                    -1,
                )
            self.nvim.funcs.nvim_buf_add_highlight(
                buf.number,
                self.highlight_namespace,
                self.options.hl.cell,
                span.end.lineno,
                0,
                span.end.colno,
            )

        if self.should_show_display_window:
            self.outputs[span].show(span.end)
        else:
            self.outputs[span].clear_interface()

    def _get_content_checksum(self) -> str:
        return hashlib.md5(
            "\n".join(self.nvim.current.buffer.api.get_lines(0, -1, True)).encode("utf-8")
        ).hexdigest()

rplugin/python3/molten/moltenbuffer.py
benlubas/molten-nvim
import os

from pynvim import Nvim
from typing import Optional, Union, List
from dataclasses import dataclass

from molten.utils import notify_error

@dataclass
class HL:
    border_norm = "MoltenOutputBorder"
    border_fail = "MoltenOutputBorderFail"
    border_succ = "MoltenOutputBorderSuccess"
    win = "MoltenOutputWin"
    win_nc = "MoltenOutputWinNC"
    foot = "MoltenOutputFooter"
    cell = "MoltenCell"

    defaults = {
        border_norm: "FloatBorder",
        border_succ: border_norm,
        border_fail: border_norm,
        win: "NormalFloat",
        win_nc: win,
        foot: "FloatFooter",
        cell: "CursorLine",
    }

class MoltenOptions:
    auto_open_output: bool
    copy_output: bool
    enter_output_behavior: str
    image_provider: str
    output_crop_border: bool
    output_show_more: bool
    output_win_border: Union[str, List[str]]
    output_win_cover_gutter: bool
    output_win_hide_on_leave: bool
    output_win_max_height: int
    output_win_max_width: int
    output_win_style: Optional[str]
    save_path: str
    use_border_highlights: bool
    wrap_output: bool
    show_mimetype_debug: bool
    nvim: Nvim
    hl: HL

    def __init__(self, nvim: Nvim):
        self.nvim = nvim
        self.hl = HL()
        # fmt: off
        CONFIG_VARS = [
            ("molten_auto_open_output", True),
            ("molten_copy_output", False),
            ("molten_enter_output_behavior", "open_then_enter"),
            ("molten_image_provider", "none"),
            ("molten_output_crop_border", True),
            ("molten_output_show_more", True),
            ("molten_output_win_border", [ "", "━", "", "" ]),
            ("molten_output_win_cover_gutter", True),
            ("molten_output_win_hide_on_leave", True),
            ("molten_output_win_max_height", 999999),
            ("molten_output_win_max_width", 999999),
            ("molten_output_win_style", False),
            ("molten_save_path", os.path.join(nvim.funcs.stdpath("data"), "molten")),
            ("molten_use_border_highlights", False),
            ("molten_wrap_output", False),
            ("molten_show_mimetype_debug", False),
        ]
        # fmt: on

        for name, default in CONFIG_VARS:
            setattr(self, name[7:], nvim.vars.get(name, default))

    def update_option(self, option: str, value):
        if option.startswith("molten_"):
            option = option[7:]
        if hasattr(self, option):
            setattr(self, option, value)
        else:
            notify_error(self.nvim, f"Invalid option passed to MoltenUpdateOption: {option}")

rplugin/python3/molten/options.py
benlubas/molten-nvim
from typing import Any, List, Optional, Union

from pynvim import Nvim
from pynvim.api import Buffer, Window

from molten.images import Canvas
from molten.outputchunks import Output, OutputStatus
from molten.options import MoltenOptions
from molten.utils import Position, notify_error


class OutputBuffer:
    nvim: Nvim
    canvas: Canvas

    output: Output

    display_buf: Buffer
    display_win: Optional[Window]

    options: MoltenOptions
    lua: Any

    def __init__(self, nvim: Nvim, canvas: Canvas, options: MoltenOptions):
        self.nvim = nvim
        self.canvas = canvas

        self.output = Output(None)

        self.display_buf = self.nvim.buffers[self.nvim.funcs.nvim_create_buf(False, True)]
        self.display_win = None

        self.options = options
        self.nvim.exec_lua("_ow = require('output_window')")
        self.lua = self.nvim.lua._ow

    def _buffer_to_window_lineno(self, lineno: int, bufno: int) -> int:
        return self.lua.calculate_window_position(bufno, lineno)

    def _get_header_text(self, output: Output) -> str:
        if output.execution_count is None:
            execution_count = "..."
        else:
            execution_count = str(output.execution_count)

        if output.status == OutputStatus.HOLD:
            status = "* On Hold"
        elif output.status == OutputStatus.DONE:
            if output.success:
                status = "✓ Done"
            else:
                status = "✗ Failed"
        elif output.status == OutputStatus.RUNNING:
            status = "... Running"
        else:
            raise ValueError("bad output.status: %s" % output.status)

        if output.old:
            old = "[OLD] "
        else:
            old = ""

        return f"{old}Out[{execution_count}]: {status}"

    def enter(self, anchor: Position) -> bool:
        entered = False
        if self.display_win is None:
            if self.options.enter_output_behavior == "open_then_enter":
                self.show(anchor)
            elif self.options.enter_output_behavior == "open_and_enter":
                self.show(anchor)
                entered = True
                self.nvim.funcs.nvim_set_current_win(self.display_win)
        elif self.options.enter_output_behavior != "no_open":
            entered = True
            self.nvim.funcs.nvim_set_current_win(self.display_win)
        if entered:
            if self.options.output_show_more:
                self.remove_window_footer()
            if self.options.output_win_hide_on_leave:
                return False
        return True

    def clear_interface(self) -> None:
        if self.display_win is not None:
            self.nvim.funcs.nvim_win_close(self.display_win, True)
            self.canvas.clear()
            self.display_win = None

    def set_win_option(self, option: str, value) -> None:
        if self.display_win:
            self.nvim.api.set_option_value(
                option,
                value,
                {"scope": "local", "win": self.display_win.handle},
            )

    def show(self, anchor: Position) -> None:
        win = self.nvim.current.window
        win_col = win.col
        win_row = self._buffer_to_window_lineno(anchor.lineno + 1, anchor.bufno)
        win_width = win.width
        win_height = win.height

        border_w, border_h = border_size(self.options.output_win_border)

        win_height -= border_h
        win_width -= border_w

        # Clear buffer:
        self.nvim.funcs.deletebufline(self.display_buf.number, 1, "$")
        # Add output chunks to buffer
        lines_str = ""
        lineno = 0
        # images are rendered with virtual lines by image.nvim
        virtual_lines = 0

        sign_col_width = 0
        text_off = self.nvim.funcs.getwininfo(win.handle)[0]["textoff"]
        if not self.options.output_win_cover_gutter:
            sign_col_width = text_off

        shape = (
            win_col + sign_col_width,
            win_row,
            win_width - sign_col_width,
            win_height,
        )
        if len(self.output.chunks) > 0:
            for chunk in self.output.chunks:
                chunktext, virt_lines = chunk.place(
                    self.display_buf.number,
                    self.options,
                    lineno,
                    shape,
                    self.canvas,
                )
                lines_str += chunktext
                lineno += chunktext.count("\n")
                virtual_lines += virt_lines

            lines = handle_progress_bars(lines_str)
            lineno = len(lines)
        else:
            lines = [lines_str]

        self.display_buf[0] = self._get_header_text(self.output)
        self.display_buf.append(lines)
        self.nvim.api.set_option_value(
            "filetype", "molten_output", {"buf": self.display_buf.handle}
        )

        # Open output window
        # assert self.display_window is None
        if win_row < win_height:
            border = self.options.output_win_border
            max_height = min(virtual_lines + lineno + 1, self.options.output_win_max_height)
            height = min(win_height - win_row, max_height)

            cropped = False
            if (
                height == win_height - win_row and max_height > height
            ):  # It didn't fit on the screen
                if self.options.output_crop_border and type(border) == list:
                    cropped = True
                    # Expand the border, so top and bottom can change independently
                    border = [border[i % len(border)] for i in range(8)]
                    border[5 % len(border)] = ""
                    height += 1

            if self.options.use_border_highlights:
                border = self.set_border_highlight(border)

            win_opts = {
                "relative": "win",
                "row": shape[1],
                "col": shape[0],
                "width": min(shape[2], self.options.output_win_max_width),
                "height": height,
                "border": border,
                "focusable": False,
            }
            if self.options.output_win_style:
                win_opts["style"] = self.options.output_win_style
            if (
                self.options.output_show_more
                and not cropped
                and height == self.options.output_win_max_height
                and len(self.display_buf) > height - border_h
            ):
                # the entire window size is shown, but the buffer still has more lines to render
                hidden_lines = len(self.display_buf) - height
                if self.options.output_win_cover_gutter and type(border) == list:
                    border_pad = border[5 % len(border)][0] * text_off
                    win_opts["footer"] = [
                        (border_pad, border[5 % len(border)][1]),
                        (f" 󰁅 {hidden_lines} More Lines ", self.options.hl.foot),
                    ]
                else:
                    win_opts["footer"] = [(f" 󰁅 {hidden_lines} More Lines ", self.options.hl.foot)]
                win_opts["footer_pos"] = "left"

            if self.display_win is None or not self.display_win.valid:  # open a new window
                self.display_win = self.nvim.api.open_win(
                    self.display_buf.number,
                    False,
                    win_opts,
                )
                hl = self.options.hl
                self.set_win_option("winhighlight", f"Normal:{hl.win},NormalNC:{hl.win_nc}")
                # TODO: Refactor once MoltenOutputWindowOpen autocommand is a thing.
                # note, the above setting will probably stay there, just so users can set highlights
                # with their other highlights
                self.set_win_option("wrap", self.options.wrap_output)
                self.set_win_option("cursorline", False)
                self.canvas.present()
            else:  # move the current window
                self.display_win.api.set_config(win_opts)

    def set_border_highlight(self, border):
        hl = self.options.hl.border_norm
        if not self.output.success:
            hl = self.options.hl.border_fail
        elif self.output.status == OutputStatus.DONE:
            hl = self.options.hl.border_succ

        if type(border) == str:
            notify_error(
                self.nvim,
                "`use_border_highlights` only works when `output_win_border` is specified as a table",
            )
            return border

        for i in range(len(border)):
            match border[i]:
                case [str(_), *_]:
                    border[i][1] = hl
                case str(_):
                    border[i] = [border[i], hl]

        return border

    def remove_window_footer(self) -> None:
        if self.display_win is not None:
            self.display_win.api.set_config({"footer": ""})


def handle_progress_bars(line_str: str) -> List[str]:
    """Progress bars like tqdm use special chars (`\\r`) and some trick to work
    This is fine for the terminal, but in a text editor we have so do some extra work
    """
    actual_lines = []
    lines = line_str.split("\n")
    for line in lines:
        parts = line.split("\r")
        last = parts[-1]
        if last != "":
            actual_lines.append(last)
            lines = actual_lines

    return actual_lines


def border_size(border: Union[str, List[str], List[List[str]]]):
    width, height = 0, 0
    match border:
        case list(b):
            height += border_char_size(1, b)
            height += border_char_size(5, b)
            width += border_char_size(7, b)
            width += border_char_size(3, b)
        case "rounded" | "single" | "double" | "solid":
            height += 2
            width += 2
        case "shadow":
            height += 1
            width += 1
    return width, height


def border_char_size(index: int, border: Union[List[str], List[List[str]]]):
    match border[index % len(border)]:
        case str(ch) | [str(ch), _]:
            return len(ch)
        case _:
            return 0

rplugin/python3/molten/outputbuffer.py
benlubas/molten-nvim
from typing import (
    Optional,
    Tuple,
    List,
    Dict,
    Any,
    Callable,
    IO,
)
from contextlib import AbstractContextManager
from enum import Enum
from abc import ABC, abstractmethod
import re

from pynvim import Nvim


from molten.images import Canvas
from molten.options import MoltenOptions
from molten.utils import notify_error


class OutputChunk(ABC):
    jupyter_data: Optional[Dict[str, Any]] = None
    jupyter_metadata: Optional[Dict[str, Any]] = None

    @abstractmethod
    def place(
        self,
        bufnr: int,
        options: MoltenOptions,
        lineno: int,
        shape: Tuple[int, int, int, int],
        canvas: Canvas,
    ) -> Tuple[str, int]:
        pass


# Adapted from [https://stackoverflow.com/a/14693789/4803382]:
ANSI_CODE_REGEX = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")


def clean_up_text(text: str) -> str:
    text = ANSI_CODE_REGEX.sub("", text)
    text = text.replace("\r\n", "\n")
    return text


class TextOutputChunk(OutputChunk):
    text: str

    def __init__(self, text: str):
        self.text = text

    def _cleanup_text(self, text: str) -> str:
        return clean_up_text(text)

    def place(
        self,
        _bufnr: int,
        options: MoltenOptions,
        _: int,
        shape: Tuple[int, int, int, int],
        _canvas: Canvas,
    ) -> Tuple[str, int]:
        text = self._cleanup_text(self.text)
        extra_lines = 0
        if options.wrap_output:  # count the number of extra lines this will need when wrapped
            win_width = shape[2]
            for line in text.split("\n"):
                if len(line) > win_width:
                    extra_lines += len(line) // win_width
        return text, extra_lines


class TextLnOutputChunk(TextOutputChunk):
    def __init__(self, text: str):
        super().__init__(text + "\n")


class BadOutputChunk(TextLnOutputChunk):
    def __init__(self, mimetypes: List[str]):
        super().__init__("<No usable MIMEtype! Received mimetypes %r>" % mimetypes)


class MimetypesOutputChunk(TextLnOutputChunk):
    def __init__(self, mimetypes: List[str]):
        super().__init__("[DEBUG] Received mimetypes: %r" % mimetypes)


class ErrorOutputChunk(TextLnOutputChunk):
    def __init__(self, name: str, message: str, traceback: List[str]):
        super().__init__(
            "\n".join(
                [
                    f"[Error] {name}: {message}",
                    "Traceback:",
                ]
                + traceback
            )
        )


class AbortedOutputChunk(TextLnOutputChunk):
    def __init__(self) -> None:
        super().__init__("<Kernel aborted with no error message.>")


class ImageOutputChunk(OutputChunk):
    def __init__(self, img_path: str):
        self.img_path = img_path

    def place(
        self,
        bufnr: int,
        _: MoltenOptions,
        lineno: int,
        _shape: Tuple[int, int, int, int],
        canvas: Canvas,
    ) -> Tuple[str, int]:
        # _x, _y, win_w, win_h = shape
        img = canvas.add_image(
            self.img_path,
            x=0,
            y=lineno + 1,
            bufnr=bufnr,
        )
        return "", canvas.img_size(img)["height"]


class OutputStatus(Enum):
    HOLD = 0
    RUNNING = 1
    DONE = 2


class Output:
    execution_count: Optional[int]
    chunks: List[OutputChunk]
    status: OutputStatus
    success: bool
    old: bool

    _should_clear: bool

    def __init__(self, execution_count: Optional[int]):
        self.execution_count = execution_count
        self.status = OutputStatus.HOLD
        self.chunks = []
        self.success = True
        self.old = False

        self._should_clear = False


def to_outputchunk(
    nvim: Nvim,
    alloc_file: Callable[
        [str, str],
        "AbstractContextManager[Tuple[str, IO[bytes]]]",
    ],
    data: Dict[str, Any],
    metadata: Dict[str, Any],
) -> OutputChunk:
    def _to_image_chunk(path: str) -> OutputChunk:
        return ImageOutputChunk(path)

    # Output chunk functions:
    def _from_image_png(imgdata: bytes) -> OutputChunk:
        import base64

        with alloc_file("png", "wb") as (path, file):
            file.write(base64.b64decode(str(imgdata)))
        return _to_image_chunk(path)

    def _from_image_svgxml(svg: str) -> OutputChunk:
        try:
            import cairosvg
            with alloc_file("png", "wb") as (path, file):
                cairosvg.svg2png(svg, write_to=file)
            return _to_image_chunk(path)
        except ImportError:
            with alloc_file("svg", "w") as (path, file):
                file.write(svg)  # type: ignore
            return _to_image_chunk(path)

    def _from_application_plotly(figure_json: Any) -> OutputChunk:
        from plotly.io import from_json
        import json

        figure = from_json(json.dumps(figure_json))

        with alloc_file("png", "wb") as (path, file):
            figure.write_image(file, engine="kaleido")
        return _to_image_chunk(path)

    def _from_latex(tex: str) -> OutputChunk:
        from pnglatex import pnglatex

        try:
            with alloc_file("png", "w") as (path, _):
                pass
            pnglatex(tex, path)
            return _to_image_chunk(path)
        except ValueError:
            notify_error(nvim, f"pnglatex was unable to render image from LaTeX: {tex}")
            return _from_plaintext(tex)

    def _from_plaintext(text: str) -> OutputChunk:
        return TextLnOutputChunk(text)

    OUTPUT_CHUNKS = {
        "image/png": _from_image_png,
        "image/svg+xml": _from_image_svgxml,
        "application/vnd.plotly.v1+json": _from_application_plotly,
        "text/latex": _from_latex,
        "text/plain": _from_plaintext,
    }

    chunk = None
    for mimetype, process_func in OUTPUT_CHUNKS.items():
        try:
            maybe_data = data.get(mimetype)
            if maybe_data is not None:
                chunk = process_func(maybe_data)  # type: ignore
                break
        except ImportError:
            continue

    if chunk is None:
        chunk = BadOutputChunk(list(data.keys()))

    chunk.jupyter_data = data
    chunk.jupyter_metadata = metadata

    return chunk

rplugin/python3/molten/outputchunks.py
benlubas/molten-nvim
from typing import Optional, Tuple, List, Dict, Generator, IO, Any
from enum import Enum
from contextlib import contextmanager
from queue import Empty as EmptyQueueException
import os
import tempfile
import json

import jupyter_client
from pynvim import Nvim

from molten.options import MoltenOptions
from molten.outputchunks import (
    Output,
    MimetypesOutputChunk,
    ErrorOutputChunk,
    TextOutputChunk,
    OutputStatus,
    to_outputchunk,
    clean_up_text,
)


class RuntimeState(Enum):
    STARTING = 0
    IDLE = 1
    RUNNING = 2


class JupyterRuntime:
    state: RuntimeState
    kernel_name: str

    kernel_manager: jupyter_client.KernelManager
    kernel_client: jupyter_client.KernelClient

    allocated_files: List[str]

    options: MoltenOptions
    nvim: Nvim

    def __init__(self, nvim: Nvim, kernel_name: str, options: MoltenOptions):
        self.state = RuntimeState.STARTING
        self.kernel_name = kernel_name
        self.nvim = nvim

        if ".json" not in self.kernel_name:
            self.external_kernel = False
            self.kernel_manager = jupyter_client.manager.KernelManager(kernel_name=kernel_name)
            self.kernel_manager.start_kernel()
            self.kernel_client = self.kernel_manager.client()
            assert isinstance(
                self.kernel_client,
                jupyter_client.blocking.client.BlockingKernelClient,
            )
            self.kernel_client.start_channels()
        else:
            kernel_file = kernel_name
            self.external_kernel = True
            # Opening JSON file
            try:
                kernel_json = json.load(open(kernel_file))
            except FileNotFoundError:
                raise ValueError(f"Could not find kernel file at path: {kernel_file}")

            # we have a kernel json
            self.kernel_manager = jupyter_client.manager.KernelManager(
                kernel_name=kernel_json["kernel_name"]
            )
            self.kernel_client = self.kernel_manager.client()
            self.kernel_client.load_connection_file(connection_file=kernel_file)

        self.allocated_files = []
        self.options = options

    def is_ready(self) -> bool:
        return self.state.value > RuntimeState.STARTING.value

    def deinit(self) -> None:
        for path in self.allocated_files:
            if os.path.exists(path):
                os.remove(path)

        if self.external_kernel is False:
            self.kernel_client.shutdown()

    def interrupt(self) -> None:
        self.kernel_manager.interrupt_kernel()

    def restart(self) -> None:
        self.state = RuntimeState.STARTING
        self.kernel_manager.restart_kernel()

    def run_code(self, code: str) -> None:
        self.kernel_client.execute(code)

    @contextmanager
    def _alloc_file(
        self, extension: str, mode: str
    ) -> Generator[Tuple[str, IO[bytes]], None, None]:
        with tempfile.NamedTemporaryFile(suffix="." + extension, mode=mode, delete=False) as file:
            path = file.name
            yield path, file
        self.allocated_files.append(path)

    def _append_chunk(self, output: Output, data: Dict[str, Any], metadata: Dict[str, Any]) -> None:
        if self.options.show_mimetype_debug:
            output.chunks.append(MimetypesOutputChunk(list(data.keys())))

        if output.success:
            output.chunks.append(to_outputchunk(self.nvim, self._alloc_file, data, metadata))

    def _tick_one(self, output: Output, message_type: str, content: Dict[str, Any]) -> bool:
        def copy_on_demand(content_ctor):
            if self.options.copy_output:
                import pyperclip

                if type(content_ctor) is str:
                    pyperclip.copy(content_ctor)
                else:
                    pyperclip.copy(content_ctor())

        if output._should_clear:
            output.chunks.clear()
            output._should_clear = False

        if message_type == "execute_input":
            output.execution_count = content["execution_count"]
            if self.external_kernel is False:
                assert output.status != OutputStatus.DONE
                if output.status == OutputStatus.HOLD:
                    output.status = OutputStatus.RUNNING
                elif output.status == OutputStatus.RUNNING:
                    output.status = OutputStatus.DONE
                else:
                    raise ValueError("bad value for output.status: %r" % output.status)
            return True
        elif message_type == "status":
            execution_state = content["execution_state"]
            assert execution_state != "starting"
            if execution_state == "idle":
                self.state = RuntimeState.IDLE
                output.status = OutputStatus.DONE
                return True
            elif execution_state == "busy":
                self.state = RuntimeState.RUNNING
                return True
            else:
                return False
        elif message_type == "execute_reply":
            # This doesn't really give us any relevant information.
            return False
        elif message_type == "execute_result":
            self._append_chunk(output, content["data"], content["metadata"])
            if "text/plain" in content["data"]:
                copy_on_demand(content["data"]["text/plain"])
            return True
        elif message_type == "error":
            output.success = False
            output.chunks.append(
                ErrorOutputChunk(content["ename"], content["evalue"], content["traceback"])
            )
            copy_on_demand(lambda: "\n\n".join(map(clean_up_text, content["traceback"])))
            return True
        elif message_type == "stream":
            copy_on_demand(content["text"])
            chunk = TextOutputChunk(content["text"])
            chunk.jupyter_data = {"text/plain": content["text"]}
            chunk.jupyter_metadata = {}
            output.chunks.append(chunk)
            return True
        elif message_type == "display_data":
            # XXX: consider content['transient'], if we end up saving execution
            # outputs.
            self._append_chunk(output, content["data"], content["metadata"])
            return True
        elif message_type == "update_display_data":
            # We don't really want to bother with this type of message.
            return False
        elif message_type == "clear_output":
            if content["wait"]:
                output._should_clear = True
            else:
                output.chunks.clear()
            return True
        # TODO: message_type == 'debug'?
        else:
            return False

    def tick(self, output: Optional[Output]) -> bool:
        did_stuff = False

        assert isinstance(
            self.kernel_client,
            jupyter_client.blocking.client.BlockingKernelClient,
        )

        if not self.is_ready():
            try:
                self.kernel_client.wait_for_ready(timeout=0)
                self.state = RuntimeState.IDLE
                did_stuff = True
            except RuntimeError:
                return False

        if output is None:
            return did_stuff

        while True:
            try:
                message = self.kernel_client.get_iopub_msg(timeout=0)

                if "content" not in message or "msg_type" not in message:
                    continue

                did_stuff_now = self._tick_one(output, message["msg_type"], message["content"])
                did_stuff = did_stuff or did_stuff_now

                if output.status == OutputStatus.DONE:
                    break
            except EmptyQueueException:
                break

        return did_stuff


def get_available_kernels() -> List[str]:
    return list(jupyter_client.kernelspec.find_kernel_specs().keys())

rplugin/python3/molten/runtime.py
benlubas/molten-nvim
from typing import Union, List

from pynvim import Nvim


class MoltenException(Exception):
    pass


def nvimui(func):  # type: ignore
    def inner(self, *args, **kwargs):  # type: ignore
        try:
            func(self, *args, **kwargs)
        except MoltenException as err:
            self.nvim.err_write("[Molten] " + str(err) + "\n")

    return inner


def _notify(nvim: Nvim, msg: str, log_level: str) -> None:
    lua = f"""
        vim.schedule_wrap(function()
            vim.notify([[[Molten] {msg}]], vim.log.levels.{log_level}, {{}})
        end)()
    """
    nvim.exec_lua(lua)


def notify_info(nvim: Nvim, msg: str) -> None:
    """Use the vim.notify API to display an info message."""
    _notify(nvim, msg, "INFO")


def notify_warn(nvim: Nvim, msg: str) -> None:
    """Use the vim.notify API to display a warning message."""
    _notify(nvim, msg, "WARN")


def notify_error(nvim: Nvim, msg: str) -> None:
    """Use the vim.notify API to display an error message."""
    _notify(nvim, msg, "ERROR")


class Position:
    bufno: int
    lineno: int
    colno: int

    def __init__(self, bufno: int, lineno: int, colno: int):
        self.bufno = bufno
        self.lineno = lineno
        self.colno = colno

    def __lt__(self, other: "Position") -> bool:
        return (self.lineno, self.colno) < (other.lineno, other.colno)

    def __le__(self, other: "Position") -> bool:
        return (self.lineno, self.colno) <= (other.lineno, other.colno)


class DynamicPosition(Position):
    nvim: Nvim
    extmark_namespace: int
    bufno: int

    extmark_id: int

    def __init__(
        self,
        nvim: Nvim,
        extmark_namespace: int,
        bufno: int,
        lineno: int,
        colno: int,
    ):
        self.nvim = nvim
        self.extmark_namespace = extmark_namespace

        self.bufno = bufno
        self.extmark_id = self.nvim.funcs.nvim_buf_set_extmark(
            self.bufno, extmark_namespace, lineno, colno, {}
        )

    def __del__(self) -> None:
        self.nvim.funcs.nvim_buf_del_extmark(self.bufno, self.extmark_namespace, self.extmark_id)

    def _get_pos(self) -> List[int]:
        out = self.nvim.funcs.nvim_buf_get_extmark_by_id(
            self.bufno, self.extmark_namespace, self.extmark_id, {}
        )
        assert isinstance(out, list) and all(isinstance(x, int) for x in out)
        return out

    @property
    def lineno(self) -> int:  # type: ignore
        return self._get_pos()[0]

    @property
    def colno(self) -> int:  # type: ignore
        return self._get_pos()[1]


class Span:
    begin: Union[Position, DynamicPosition]
    end: Union[Position, DynamicPosition]
    bufno: int

    def __init__(
        self,
        begin: Union[Position, DynamicPosition],
        end: Union[Position, DynamicPosition],
    ):
        self.begin = begin
        self.end = end
        assert self.begin.bufno == self.end.bufno

    def __contains__(self, pos: Union[Position, DynamicPosition]) -> bool:
        return self.begin.bufno == pos.bufno and self.begin <= pos and pos < self.end

    def get_text(self, nvim: Nvim) -> str:
        assert self.begin.bufno == self.end.bufno
        bufno = self.begin.bufno

        lines: List[str] = nvim.funcs.nvim_buf_get_lines(
            bufno, self.begin.lineno, self.end.lineno + 1, True
        )

        if len(lines) == 1:
            return lines[0][self.begin.colno : self.end.colno]
        else:
            return "\n".join(
                [lines[0][self.begin.colno :]] + lines[1:-1] + [lines[-1][: self.end.colno]]
            )

rplugin/python3/molten/utils.py
benlubas/molten-nvim