From d93164fed2ddcb2d4ebb15322f8592d56104f3ff Mon Sep 17 00:00:00 2001
From: ESSO0428 <92996726+ESSO0428@users.noreply.github.com>
Date: Tue, 16 Jul 2024 08:13:48 +0800
Subject: [PATCH] Enhance Parser to Support # %% with Additional Comments
 (#127)

* refactor(parser): support # %% with additional comments
* fix(parser): import re for regex matching
* refactor(parser): support # %% with additional comments
* fix: correct cell separator detection for magic commands
* docs: add comments for '# %% Cell Title' feature
* test: add tests for '# %% Cell Title' in test_buffer.py
---
 lua/jupynium/cells.lua |  5 ++++-
 src/jupynium/buffer.py |  5 ++++-
 tests/test_buffer.py   | 16 ++++++++++++++++
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/lua/jupynium/cells.lua b/lua/jupynium/cells.lua
index 7105e49..a6623f5 100644
--- a/lua/jupynium/cells.lua
+++ b/lua/jupynium/cells.lua
@@ -12,7 +12,10 @@ function M.line_type(line)
     return "empty"
   elseif vim.startswith(line, "# %% [md]") or vim.startswith(line, "# %% [markdown]") then
     return "cell separator: markdown"
-  elseif vim.fn.trim(line) == "# %%" then
+  -- Treat '# %% Cell Title' as a code cell
+  -- But not magic commands such as '# %%timeit'.
+  -- See: https://github.com/kiyoon/jupynium.nvim/pull/127
+  elseif vim.fn.trim(string.sub(line, 1, 5)) == "# %%" then
     return "cell separator: code"
   elseif vim.startswith(line, "# ---") then
     return "metadata"
diff --git a/src/jupynium/buffer.py b/src/jupynium/buffer.py
index bf68b0a..93008c5 100644
--- a/src/jupynium/buffer.py
+++ b/src/jupynium/buffer.py
@@ -75,7 +75,10 @@ def full_analyse_buf(self, header_cell_type: str = "header"):
                 num_rows_per_cell.append(num_rows_this_cell)
                 num_rows_this_cell = 1
                 cell_types.append("markdown")
-            elif line.strip() == "# %%":
+            # Treat '# %% Cell Title' as a code cell
+            # But not magic commands such as '# %%timeit'.
+            # See: https://github.com/kiyoon/jupynium.nvim/pull/127
+            elif line[:5].strip() == "# %%":
                 num_rows_per_cell.append(num_rows_this_cell)
                 num_rows_this_cell = 1
                 cell_types.append("code")
diff --git a/tests/test_buffer.py b/tests/test_buffer.py
index 57276cf..cf60809 100644
--- a/tests/test_buffer.py
+++ b/tests/test_buffer.py
@@ -11,6 +11,12 @@ def test_buffer_1():
     assert buffer.cell_types == ["header", "code"]
 
 
+def test_buffer_with_cell_title_1():
+    buffer = JupyniumBuffer(["a", "b", "c", "# %% Cell Title", "d", "e", "f"])
+    assert buffer.num_rows_per_cell == [3, 4]
+    assert buffer.cell_types == ["header", "code"]
+
+
 def test_magic_command_1():
     """
     Everything else except magic commands should be preserved after __init__() or fully_analysed_buf().
@@ -21,6 +27,16 @@ def test_magic_command_1():
     assert code_cell_content == "%time\ne\nf"
 
 
+def test_double_magic_command_1():
+    """
+    Everything else except double magic commands should be preserved after __init__() or fully_analyse_buf().
+    """
+    lines = ["a", "b", "c", "# %% Cell Title", "# %%timeit", "e", "f"]
+    buffer = JupyniumBuffer(lines)
+    code_cell_content = buffer.get_cell_text(1)
+    assert code_cell_content == "%%timeit\ne\nf"
+
+
 def test_buffer_markdown():
     buffer = JupyniumBuffer(["a", "b", "c", "# %% [md]", "d", "# %%", "f"])
     assert buffer.num_rows_per_cell == [3, 2, 2]