fix: hide images when appropriate (#44)

rplugin/python3/molten/__init__.py

@@ -120,8 +120,10 @@ def _clear_interface(self) -> None:
             for molten_kernel in molten_kernels:
                 molten_kernel.clear_interface()
                 molten_kernel.clear_open_output_windows()
+                molten_kernel.clear_virt_outputs()
 
         assert self.canvas is not None
+        self.canvas.clear()
         self.canvas.present()
 
     def _update_interface(self) -> None:

rplugin/python3/molten/images.py

@@ -56,6 +56,7 @@ def add_image(
     ) -> str:
         """
         Add an image to the canvas.
+        Takes effect after a call to present()
 
         Parameters
         - path: str
@@ -73,6 +74,17 @@ def add_image(
         str the identifier for the image
         """
 
+    @abstractmethod
+    def remove_image(self, identifier: str) -> None:
+        """
+        Remove an image from the canvas. In practice this is just hiding the image
+        Takes effect after a call to present()
+
+        Parameters
+        - identifier: str
+          The identifier for the image to remove.
+        """
+
 
 class NoCanvas(Canvas):
     def __init__(self) -> None:
@@ -102,6 +114,9 @@ def add_image(
     ) -> None:
         pass
 
+    def remove_image(self, _identifier: str) -> None:
+        pass
+
 
 class ImageNvimCanvas(Canvas):
     nvim: Nvim
@@ -173,6 +188,9 @@ def add_image(
             return img
         return path
 
+    def remove_image(self, identifier: str) -> None:
+        self.to_make_invisible.add(identifier)
+
 
 def get_canvas_given_provider(name: str, nvim: Nvim) -> Canvas:
     if name == "none":

rplugin/python3/molten/moltenbuffer.py

@@ -91,6 +91,7 @@ def restart(self, delete_outputs: bool = False) -> None:
             self.outputs = {}
             self.clear_interface()
             self.clear_open_output_windows()
+            self.clear_virt_outputs()
 
         self.runtime.restart()
 
@@ -173,6 +174,10 @@ def clear_open_output_windows(self) -> None:
         for output in self.outputs.values():
             output.clear_float_win()
 
+    def clear_virt_outputs(self) -> None:
+        for cell, output in self.outputs.items():
+            output.clear_virt_output(cell.bufno)
+
     def _get_selected_span(self) -> Optional[CodeCell]:
         current_position = self._get_cursor_position()
         selected = None
@@ -200,6 +205,7 @@ def delete_cell(self) -> None:
             return
 
         self.outputs[self.selected_cell].clear_float_win()
+        self.outputs[self.selected_cell].clear_virt_output(self.selected_cell.bufno)
         self.selected_cell.clear_interface(self.highlight_namespace)
         del self.outputs[self.selected_cell]
         self.selected_cell = None

rplugin/python3/molten/outputbuffer.py

@@ -4,7 +4,7 @@
 from pynvim.api import Buffer, Window
 
 from molten.images import Canvas
-from molten.outputchunks import Output, OutputStatus
+from molten.outputchunks import ImageOutputChunk, Output, OutputStatus
 from molten.options import MoltenOptions
 from molten.position import DynamicPosition, Position
 from molten.utils import notify_error
@@ -93,7 +93,6 @@ def enter(self, anchor: Position) -> bool:
     def clear_float_win(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
         if self.display_virt_lines is not None:
             del self.display_virt_lines
@@ -102,6 +101,14 @@ def clear_float_win(self) -> None:
     def clear_virt_output(self, bufnr: int) -> None:
         if self.virt_text_id is not None:
             self.nvim.funcs.nvim_buf_del_extmark(bufnr, self.extmark_namespace, self.virt_text_id)
+        # clear the image too
+        redraw = False
+        for chunk in self.output.chunks:
+            if isinstance(chunk, ImageOutputChunk) and chunk.img_identifier is not None:
+                self.canvas.remove_image(chunk.img_identifier)
+                redraw = True
+        if redraw:
+            self.canvas.present()
 
     def set_win_option(self, option: str, value) -> None:
         if self.display_win:
@@ -189,7 +196,7 @@ def show_floating_win(self, anchor: Position) -> None:
         win = self.nvim.current.window
         win_col = win.col
         win_row = self._buffer_to_window_lineno(anchor.lineno + 1)
-        if win_row == 0: # anchor position is off screen
+        if win_row == 0:  # anchor position is off screen
             return
         win_width = win.width
         win_height = win.height

rplugin/python3/molten/outputchunks.py

@@ -126,6 +126,7 @@ class ImageOutputChunk(OutputChunk):
     def __init__(self, img_path: str):
         self.img_path = img_path
         self.output_type = "display_data"
+        self.img_identifier = None
 
     def place(
         self,
@@ -136,13 +137,13 @@ def place(
         canvas: Canvas,
         _hard_wrap: bool,
     ) -> Tuple[str, int]:
-        img = canvas.add_image(
+        self.img_identifier = canvas.add_image(
             self.img_path,
             x=0,
             y=lineno + 1,
             bufnr=bufnr,
         )
-        return "", canvas.img_size(img)["height"]
+        return "", canvas.img_size(self.img_identifier)["height"]
 
 
 class OutputStatus(Enum):