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]