Add option to share Magma between buffers

rplugin/python3/magma/__init__.py

@@ -28,13 +28,16 @@ class Magma:
 
     options: MagmaOptions
 
+    magma_buffers: Dict[str, MagmaBuffer]
+
     def __init__(self, nvim: Nvim):
         self.nvim = nvim
         self.initialized = False
 
         self.canvas = None
         self.buffers = {}
         self.timer = None
+        self.magma_buffers = {}
 
     def _initialize(self) -> None:
         assert not self.initialized
@@ -136,8 +139,16 @@ def _ask_for_choice(
         else:
             return options[index - 1]
 
-    def _initialize_buffer(self, kernel_name: str) -> MagmaBuffer:
+    def _initialize_buffer(self, kernel_name: str, shared: bool) -> MagmaBuffer:
         assert self.canvas is not None
+        if shared:
+            magma = self.magma_buffers.get(kernel_name)
+            if magma is not None:
+                magma.add_nvim_buffer(self.nvim.current.buffer)
+                self.buffers[self.nvim.current.buffer.number] = magma
+
+                return magma
+
         magma = MagmaBuffer(
             self.nvim,
             self.canvas,
@@ -148,18 +159,26 @@ def _initialize_buffer(self, kernel_name: str) -> MagmaBuffer:
             kernel_name,
         )
 
+        # Can only share one instance of a kernel
+        self.magma_buffers[kernel_name] = magma
         self.buffers[self.nvim.current.buffer.number] = magma
 
         return magma
 
-    @pynvim.command("MagmaInit", nargs="?", sync=True, complete='file')  # type: ignore
+    @pynvim.command("MagmaInit", nargs="*", sync=True, complete='file')  # type: ignore
     @nvimui  # type: ignore
     def command_init(self, args: List[str]) -> None:
         self._initialize_if_necessary()
 
         if args:
+            args = args[0].split(" ")
             kernel_name = args[0]
-            self._initialize_buffer(kernel_name)
+            shared = False
+            if len(args) > 1:
+                shared = True
+            print(args)
+
+            self._initialize_buffer(kernel_name, shared)
         else:
             PROMPT = "Select the kernel to launch:"
             available_kernels = get_available_kernels()
@@ -384,7 +403,7 @@ def command_load(self, args: List[str]) -> None:
             MagmaIOError.assert_has_key(data, "kernel", str)
             kernel_name = data["kernel"]
 
-            magma = self._initialize_buffer(kernel_name)
+            magma = self._initialize_buffer(kernel_name, shared=False)
 
             load(magma, data)
 

rplugin/python3/magma/magmabuffer.py

@@ -1,4 +1,4 @@
-from typing import Optional, Dict
+from typing import Optional, Dict, List
 from queue import Queue
 import hashlib
 
@@ -19,7 +19,7 @@ class MagmaBuffer:
     canvas: Canvas
     highlight_namespace: int
     extmark_namespace: int
-    buffer: Buffer
+    buffers: List[Buffer]
 
     runtime: JupyterRuntime
 
@@ -47,7 +47,7 @@ def __init__(
         self.canvas = canvas
         self.highlight_namespace = highlight_namespace
         self.extmark_namespace = extmark_namespace
-        self.buffer = buffer
+        self.buffers = [buffer]
 
         self._doautocmd("MagmaInitPre")
 
@@ -65,6 +65,9 @@ def __init__(
 
         self._doautocmd("MagmaInitPost")
 
+    def add_nvim_buffer(self, buffer: Buffer) -> None:
+        self.buffers.append(buffer)
+
     def _doautocmd(self, autocmd: str) -> None:
         assert " " not in autocmd
         self.nvim.command(f"doautocmd User {autocmd}")
@@ -153,12 +156,13 @@ def clear_interface(self) -> None:
         if self.updating_interface:
             return
 
-        self.nvim.funcs.nvim_buf_clear_namespace(
-            self.buffer.number,
-            self.highlight_namespace,
-            0,
-            -1,
-        )
+        for buffer in self.buffers:
+            self.nvim.funcs.nvim_buf_clear_namespace(
+                buffer.number,
+                self.highlight_namespace,
+                0,
+                -1,
+            )
         # and self.nvim.funcs.winbufnr(self.display_window) != -1:
         if self.selected_cell is not None and self.selected_cell in self.outputs:
             self.outputs[self.selected_cell].clear_interface()
@@ -194,9 +198,12 @@ def delete_cell(self) -> None:
         del self.outputs[self.selected_cell]
 
     def update_interface(self) -> None:
-        if self.buffer.number != self.nvim.current.buffer.number:
+        buffer_numbers = [
+            buffer.number for buffer in self.buffers
+        ]
+        if self.nvim.current.buffer.number not in buffer_numbers:
             return
-        if self.buffer.number != self.nvim.current.window.buffer.number:
+        if self.nvim.current.window.buffer.number not in buffer_numbers:
             return
 
         self.clear_interface()
@@ -220,41 +227,42 @@ def update_interface(self) -> None:
         self.updating_interface = False
 
     def _show_selected(self, span: Span) -> None:
-        if span.begin.lineno == span.end.lineno:
-            self.nvim.funcs.nvim_buf_add_highlight(
-                self.buffer.number,
-                self.highlight_namespace,
-                self.options.cell_highlight_group,
-                span.begin.lineno,
-                span.begin.colno,
-                span.end.colno,
-            )
-        else:
-            self.nvim.funcs.nvim_buf_add_highlight(
-                self.buffer.number,
-                self.highlight_namespace,
-                self.options.cell_highlight_group,
-                span.begin.lineno,
-                span.begin.colno,
-                -1,
-            )
-            for lineno in range(span.begin.lineno + 1, span.end.lineno):
+        for buffer in self.buffers:
+            if span.begin.lineno == span.end.lineno:
                 self.nvim.funcs.nvim_buf_add_highlight(
-                    self.buffer.number,
+                    buffer.number,
                     self.highlight_namespace,
                     self.options.cell_highlight_group,
-                    lineno,
-                    0,
+                    span.begin.lineno,
+                    span.begin.colno,
+                    span.end.colno,
+                )
+            else:
+                self.nvim.funcs.nvim_buf_add_highlight(
+                    buffer.number,
+                    self.highlight_namespace,
+                    self.options.cell_highlight_group,
+                    span.begin.lineno,
+                    span.begin.colno,
                     -1,
                 )
-            self.nvim.funcs.nvim_buf_add_highlight(
-                self.buffer.number,
-                self.highlight_namespace,
-                self.options.cell_highlight_group,
-                span.end.lineno,
-                0,
-                span.end.colno,
-            )
+                for lineno in range(span.begin.lineno + 1, span.end.lineno):
+                    self.nvim.funcs.nvim_buf_add_highlight(
+                        buffer.number,
+                        self.highlight_namespace,
+                        self.options.cell_highlight_group,
+                        lineno,
+                        0,
+                        -1,
+                    )
+                self.nvim.funcs.nvim_buf_add_highlight(
+                    buffer.number,
+                    self.highlight_namespace,
+                    self.options.cell_highlight_group,
+                    span.end.lineno,
+                    0,
+                    span.end.colno,
+                )
 
         if self.should_open_display_window:
             self.outputs[span].show(span.end)