From bee4bb5b43d20d2cdd88f1555af92d362f3aa362 Mon Sep 17 00:00:00 2001
From: Ben Lubas <56943754+benlubas@users.noreply.github.com>
Date: Sat, 6 Jan 2024 10:48:53 -0500
Subject: [PATCH] fix: progress bar exports and similar (#122)

---
 rplugin/python3/molten/outputbuffer.py | 25 +++++++------------------
 rplugin/python3/molten/outputchunks.py | 21 +++++++++++++++++++--
 rplugin/python3/molten/runtime.py      | 12 +++++-------
 3 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/rplugin/python3/molten/outputbuffer.py b/rplugin/python3/molten/outputbuffer.py
index 9a86911..679ca87 100644
--- a/rplugin/python3/molten/outputbuffer.py
+++ b/rplugin/python3/molten/outputbuffer.py
@@ -144,7 +144,7 @@ def build_output_text(self, shape, buf: int, virtual: bool) -> Tuple[List[str],
                     shape,
                     self.canvas,
                     virtual,
-                    winnr = self.nvim.current.window.handle if virtual else None,
+                    winnr=self.nvim.current.window.handle if virtual else None,
                 )
                 lines_str += chunktext
                 lineno += chunktext.count("\n")
@@ -156,13 +156,17 @@ def build_output_text(self, shape, buf: int, virtual: bool) -> Tuple[List[str],
                 lines_str = lines_str[:limit]
                 lines_str += f"\n...truncated to {limit} chars\n"
 
-            lines = handle_progress_bars(lines_str)
+            lines = lines_str.split("\n")
             lineno = len(lines) + virtual_lines
         else:
             lines = []
 
+        # Remove trailing empty lines
+        while len(lines) > 0 and lines[-1] == "":
+            lines.pop()
+
         lines.insert(0, self._get_header_text(self.output))
-        return lines, lineno
+        return lines, len(lines) - 1
 
     def show_virtual_output(self, anchor: Position) -> None:
         if self.displayed_status == OutputStatus.DONE and self.virt_text_id is not None:
@@ -359,21 +363,6 @@ def remove_window_footer(self) -> 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)
-
-    return actual_lines
-
-
 def border_size(border: Union[str, List[str], List[List[str]]]):
     width, height = 0, 0
     match border:
diff --git a/rplugin/python3/molten/outputchunks.py b/rplugin/python3/molten/outputchunks.py
index 653d564..0c398ab 100644
--- a/rplugin/python3/molten/outputchunks.py
+++ b/rplugin/python3/molten/outputchunks.py
@@ -59,6 +59,9 @@ def __init__(self, text: str):
         self.text = text
         self.output_type = "display_data"
 
+    def __repr__(self) -> str:
+        return f"TextOutputChunk(\"{self.text}\")"
+
     def place(
         self,
         _bufnr: int,
@@ -89,8 +92,8 @@ def place(
                     for _ in range(len(line) // win_width):
                         splits.append(line[index * win_width : (index + 1) * win_width])
                         index += 1
-                    else:
-                        splits.append(line[index * win_width :])
+                    splits.append(line[index * win_width :])
+
                 lines.extend(splits)
                 text = "\n".join(lines)
             else:
@@ -187,6 +190,20 @@ def __init__(self, execution_count: Optional[int]):
 
         self._should_clear = False
 
+    def merge_text_chunks(self):
+        """Merge the last two chunks if they are text chunks, and text on a line before \r
+        character, this is b/c outputs before a \r aren't shown, and so, should be deleted"""
+        if (
+            len(self.chunks) >= 2
+                and isinstance((c1 := self.chunks[-2]), TextOutputChunk)
+                and isinstance((c2 := self.chunks[-1]), TextOutputChunk)
+        ):
+            c1.text += c2.text
+            c1.text = "\n".join([re.sub(r".*\r", "", x) for x in c1.text.split("\n")[:-1]])
+            c1.jupyter_data = { "text/plain": c1.text }
+            self.chunks.pop()
+        elif len(self.chunks) > 0 and isinstance((c1 := self.chunks[0]), TextOutputChunk):
+            c1.text = "\n".join([re.sub(r".*\r", "", x) for x in c1.text.split("\n")[:-1]])
 
 def to_outputchunk(
     nvim: Nvim,
diff --git a/rplugin/python3/molten/runtime.py b/rplugin/python3/molten/runtime.py
index 3b7227d..4ca9d5f 100644
--- a/rplugin/python3/molten/runtime.py
+++ b/rplugin/python3/molten/runtime.py
@@ -108,9 +108,10 @@ def _append_chunk(self, output: Output, data: Dict[str, Any], metadata: Dict[str
             output.chunks.append(MimetypesOutputChunk(list(data.keys())))
 
         if output.success:
-            output.chunks.append(
-                to_outputchunk(self.nvim, self._alloc_file, data, metadata, self.options)
-            )
+            chunk = to_outputchunk(self.nvim, self._alloc_file, data, metadata, self.options)
+            output.chunks.append(chunk)
+            if isinstance(chunk, TextOutputChunk) and chunk.text.startswith("\r"):
+                output.merge_text_chunks()
 
     def _tick_one(self, output: Output, message_type: str, content: Dict[str, Any]) -> bool:
         def copy_on_demand(content_ctor):
@@ -167,10 +168,7 @@ def copy_on_demand(content_ctor):
             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)
+            self._append_chunk(output, {"text/plain": content["text"]}, {})
             return True
         elif message_type == "display_data":
             # XXX: consider content['transient'], if we end up saving execution