Multi Buffer Support (#4)

README.md

@@ -55,9 +55,9 @@ Jupyter provides a rich set of outputs. To see what we can currently handle, see
 
 A list of the commands and their arguments. Args in `[]` are optional
 
-| Command                  | Arguments             | Description                        |
-|--------------------------|-----------------------|------------------------------------|
-| `MoltenInit`              | `[kernel]`            | Initialize a kernel for the current buffer. If no kernel is given, prompts the user |
+| Command                   | Arguments             | Description                        |
+|---------------------------|-----------------------|------------------------------------|
+| `MoltenInit`              | `["shared"] [kernel]` | Initialize a kernel for the current buffer. If `shared` is passed as the first value, this buffer will use an already running kernel. If no kernel is given, prompts the user. |
 | `MoltenDeinit`            | none                  | De-initialize the current buffer's runtime and molten instance. (called automatically on vim close/buffer unload) |
 | `MoltenEvaluateLine`      | none                  | Evaluate the current line |
 | `MoltenEvaluateVisual`    | none                  | Evaluate the visual selection (**cannot be called with a range!**) |
@@ -70,14 +70,14 @@ A list of the commands and their arguments. Args in `[]` are optional
 | `MoltenInterrupt`         | none                  | Sends a keyboard interrupt to the kernel which stops any currently running code. (does nothing if there's no current output) |
 | `MoltenRestart`           | `[!]`                 | Shuts down a restarts the current kernel. Deletes all outputs if used with a bang |
 | `MoltenSave`              | `[path]`              | Save the current cells and evaluated outputs into a JSON file. When path is specified, save the file to `path`, otherwise save to `g:molten_save_path` |
-| `MoltenLoad`              | `[path]`              | Loads cell locations and output from a JSON file generated by `MoltenSave`. path functions the same as `MoltenSave` |
+| `MoltenLoad`              | `["shared"] [path]`   | Loads cell locations and output from a JSON file generated by `MoltenSave`. path functions the same as `MoltenSave`. If `shared` is specified, the buffer shares an already running kernel. |
 
 ## Keybindings
 
-TODO: wiki link
 The commands above should be mapped to keys for the best experience. There are more detailed setups
-in the Wiki, but here are some example bindings. Pay attention to `MoltenEvaluateVisual` and
-`MoltenEnterOutput`, as they require a little special attention
+in the [Wiki](https://github.com/benlubas/molten-nvim/wiki), but here are some example bindings.
+Pay attention to `MoltenEvaluateVisual` and `MoltenEnterOutput`, as they need to be run in...odd
+ways.
 
 ```lua
 vim.keymap.set("n", "<localleader>R", ":MoltenEvaluateOperator<CR>",
@@ -237,10 +237,10 @@ In the Jupyter protocol, most output-related messages provide a dictionary of mi
 Here is a list of the currently handled mime-types:
 
 - `text/plain`: Plain text. Shown as text in the output window's buffer.
-- `image/png`: A PNG image. Shown according to `g:molten_image_provider`.
-- `image/svg+xml`: A SVG image. Rendered into a PNG with [CairoSVG](https://cairosvg.org/) and shown with [Image.nvim](https://github.com/3rd/image.nvim).
-- `application/vnd.plotly.v1+json`: A Plotly figure. Rendered into a PNG with [Plotly](https://plotly.com/python/) + [Kaleido](https://github.com/plotly/Kaleido) and shown with [Image.nvim](https://github.com/3rd/image.nvim).
-- `text/latex`: A LaTeX formula. Rendered into a PNG with [pnglatex](https://pypi.org/project/pnglatex/) and shown with [Image.nvim](https://github.com/3rd/image.nvim).
+- `image/png`: A PNG image.
+- `image/svg+xml`: A SVG image. Rendered into a PNG with [CairoSVG](https://cairosvg.org/).
+- `application/vnd.plotly.v1+json`: A Plotly figure. Rendered into a PNG with [Plotly](https://plotly.com/python/) + [Kaleido](https://github.com/plotly/Kaleido)
+- `text/latex`: A LaTeX formula. Rendered into a PNG with [pnglatex](https://pypi.org/project/pnglatex/)
 
 This already provides quite a bit of basic functionality, but if you find a use case for a mime-type that isn't currently supported, feel free to open an issue and/or PR!
 

rplugin/python3/molten/__init__.py

@@ -9,7 +9,7 @@
 from molten.options import MoltenOptions
 from molten.outputbuffer import OutputBuffer
 from molten.runtime import get_available_kernels
-from molten.utils import DynamicPosition, MoltenException, Span, nvimui
+from molten.utils import DynamicPosition, MoltenException, Span, notify_error, notify_info, nvimui
 from pynvim import Nvim
 
 
@@ -22,19 +22,23 @@ class Molten:
     highlight_namespace: int
     extmark_namespace: int
 
-    buffers: Dict[int, MoltenBuffer]
-
     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
@@ -79,10 +83,7 @@ def _initialize_if_necessary(self) -> None:
     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 <kernel_name>` to \
-                initialize."
-            )
+            raise MoltenException("Molten is not initialized; run `:MoltenInit` to initialize.")
         return maybe_molten
 
     def _clear_interface(self) -> None:
@@ -123,8 +124,19 @@ def _ask_for_choice(self, preface: str, options: List[str]) -> Optional[str]:
         else:
             return options[index - 1]
 
-    def _initialize_buffer(self, kernel_name: str) -> MoltenBuffer:
+    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,
@@ -135,51 +147,66 @@ def _initialize_buffer(self, kernel_name: str) -> MoltenBuffer:
             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
+    @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)
+            self._initialize_buffer(kernel_name, shared=shared)
         else:
             PROMPT = "Select the kernel to launch:"
             available_kernels = get_available_kernels()
-            if self.nvim.exec_lua("return vim.ui.select ~= nil"):
-                self.nvim.exec_lua(
-                    """
-                        vim.ui.select(
-                            {%s},
-                            {prompt = "%s"},
-                            function(choice)
-                                if choice ~= nil then
-                                    vim.cmd("MoltenInit " .. choice)
-                                end
-                            end
-                        )
-                    """
-                    % (
-                        ", ".join(repr(x) for x in available_kernels),
-                        PROMPT,
-                    )
+
+            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)
                 )
-            else:
-                kernel_name = self._ask_for_choice(
-                    PROMPT,
-                    available_kernels,  # type: ignore
+
+            if len(available_kernels) == 0:
+                notify_error(
+                    self.nvim,
+                    "Molten has no running kernels to share. Please use :MoltenInit without the shared option.",
                 )
-                if kernel_name is not None:
-                    self.command_init([kernel_name])
+                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
+                                print("\\n")
+                                vim.cmd("MoltenInit {'shared ' if shared else ''}" .. choice)
+                            end
+                        end
+                    )
+                end)()
+            """
+            self.nvim.exec_lua(lua, async_=False)
 
     def _deinit_buffer(self, molten: MoltenBuffer) -> None:
         molten.deinit()
-        del self.buffers[molten.buffer.number]
+        for buf in molten.buffers:
+            del self.buffers[buf.number]
 
     @pynvim.command("MoltenDeinit", nargs=0, sync=True)  # type: ignore
     @nvimui  # type: ignore
@@ -231,10 +258,9 @@ def function_update_option(self, args) -> None:
             option, value = args
             molten.options.update_option(option, value)
         else:
-            self.nvim.api.notify(
+            notify_error(
+                self.nvim,
                 f"MoltenUpdateOption: wrong number of arguments, expected 2, given {len(args)}",
-                pynvim.logging.ERROR,
-                {"title": "Molten"},
             )
 
     @pynvim.command("MoltenEnterOutput", sync=True)  # type: ignore
@@ -357,10 +383,11 @@ def command_hide_output(self) -> None:
     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, self.nvim.current.buffer)
+            path = get_default_save_file(self.options, buf)
 
         dirname = os.path.dirname(path)
         if not os.path.exists(dirname):
@@ -370,20 +397,27 @@ def command_save(self, args: List[str]) -> None:
         assert molten is not None
 
         with open(path, "w") as file:
-            json.dump(save(molten), file)
+            json.dump(save(molten, buf.number), file)
 
-    @pynvim.command("MoltenLoad", nargs="?", sync=True)  # type: ignore
+    @pynvim.command("MoltenLoad", nargs="*", sync=True)  # type: ignore
     @nvimui  # type: ignore
     def command_load(self, args: List[str]) -> None:
         self._initialize_if_necessary()
 
-        if args:
+        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; MoltenLoad initializes Molten.")
+            raise MoltenException(
+                "Molten is already initialized for this buffer; MoltenLoad initializes Molten."
+            )
 
         with open(path) as file:
             data = json.load(file)
@@ -398,9 +432,9 @@ def command_load(self, args: List[str]) -> None:
             MoltenIOError.assert_has_key(data, "kernel", str)
             kernel_name = data["kernel"]
 
-            molten = self._initialize_buffer(kernel_name)
+            molten = self._initialize_buffer(kernel_name, shared=shared)
 
-            load(molten, data)
+            load(molten, self.nvim.current.buffer, data)
 
             self._update_interface()
         except MoltenIOError as err:

rplugin/python3/molten/io.py

@@ -36,7 +36,7 @@ def get_default_save_file(options: MoltenOptions, buffer: Buffer) -> str:
     return os.path.join(options.save_path, mangled_name + ".json")
 
 
-def load(moltenbuffer: MoltenBuffer, data: Dict[str, Any]) -> None:
+def load(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"]:
@@ -54,14 +54,14 @@ def load(moltenbuffer: MoltenBuffer, data: Dict[str, Any]) -> None:
         begin_position = DynamicPosition(
             moltenbuffer.nvim,
             moltenbuffer.extmark_namespace,
-            moltenbuffer.buffer.number,
+            nvim_buffer.number,
             cell["span"]["begin"]["lineno"],
             cell["span"]["begin"]["colno"],
         )
         end_position = DynamicPosition(
             moltenbuffer.nvim,
             moltenbuffer.extmark_namespace,
-            moltenbuffer.buffer.number,
+            nvim_buffer.number,
             cell["span"]["end"]["lineno"],
             cell["span"]["end"]["colno"],
         )
@@ -97,7 +97,8 @@ def load(moltenbuffer: MoltenBuffer, data: Dict[str, Any]) -> None:
         )
 
 
-def save(moltenbuffer: MoltenBuffer) -> Dict[str, Any]:
+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,
@@ -127,6 +128,6 @@ def save(moltenbuffer: MoltenBuffer) -> Dict[str, Any]:
                     and chunk.jupyter_metadata is not None
                 ],
             }
-            for span, output in moltenbuffer.outputs.items()
+            for span, output in moltenbuffer.outputs.items() if span.begin.bufno == nvim_buffer
         ],
     }