fix: derive notebook cell boundaries from the AST, not a line prefix

autohands/add_notebook_quotes.py

@@ -4,11 +4,81 @@
 ./add_notebook_quotes.py /path/to/input /path/to/output
 """
 
-from typing import Iterable, List
+import ast
+
+from typing import Iterable, List, Tuple
 
 from sys import argv
 
 
+def _narrative_docstring_ranges(lines: List[str]) -> List[Tuple[int, int]]:
+    """Locate the narrative docstring blocks of a script, by parsing it.
+
+    Returns ``(start, end)`` 0-based line-index pairs, one per block, in source
+    order — ``start`` is the opening-delimiter line, ``end`` the closing one.
+
+    A narrative docstring is a **bare string expression statement at module
+    level, written at column 0 with a triple-quote delimiter**. That is the shape
+    the notebook generator turns into a markdown cell. Deriving it from the
+    parsed source is what separates it from a string *bound to a name*::
+
+        s = '''
+        literal
+        '''
+
+    whose closing delimiter also sits at column 0. A line-prefix test cannot see
+    that opener — it is indented behind ``s = `` — but does match the closer, so
+    it reads the literal's end as a docstring boundary and inverts every cell
+    boundary after it: the enclosing code cell becomes a ``SyntaxError`` and the
+    code that follows is emitted as prose. ``ast`` tells the two apart by node
+    type, so the confusion cannot arise.
+
+    Two shapes raise rather than convert, because a wrong guess here silently
+    ships a broken notebook — the same reasoning as the stray-``# %%`` guard in
+    ``add_notebook_quotes``:
+
+    * a script that does not parse, since there is then no parsed source to
+      derive cell boundaries from;
+    * a column-0 single-line docstring, whose opening and closing delimiters
+      share a line and so cannot bracket a cell. Write them on their own lines.
+    """
+    try:
+        module = ast.parse("".join(lines))
+    except SyntaxError as exc:
+        where = f"line {exc.lineno}" if exc.lineno else "unknown line"
+        raise ValueError(
+            f"source script does not parse as Python ({where}: {exc.msg}) — "
+            f"notebook cell boundaries are derived from the parsed source, so a "
+            f"script that does not compile cannot be converted."
+        ) from exc
+
+    ranges: List[Tuple[int, int]] = []
+    for node in module.body:
+        if not isinstance(node, ast.Expr):
+            continue
+        value = node.value
+        if not (isinstance(value, ast.Constant) and isinstance(value.value, str)):
+            continue
+        if node.col_offset != 0:
+            continue
+
+        start = node.lineno - 1
+        end = node.end_lineno - 1
+        if not (lines[start].startswith('"""') or lines[start].startswith("'''")):
+            # A column-0 string statement written with a single-quote delimiter
+            # was never a cell boundary; leave it as code, as it always was.
+            continue
+        if start == end:
+            raise ValueError(
+                f"line {start + 1} is a single-line docstring "
+                f"({lines[start].strip()!r}) — its opening and closing "
+                f"delimiters share a line, so it cannot bracket a notebook "
+                f"cell. Write the delimiters on their own lines."
+            )
+        ranges.append((start, end))
+    return ranges
+
+
 def strip_env_declarations(lines: List[str]) -> List[str]:
     """Remove in-file env declarations before notebook / markdown conversion.
 
@@ -27,13 +97,20 @@ def strip_env_declarations(lines: List[str]) -> List[str]:
       removed at runtime (it now raises in ``read_env_declaration``), but a stray
       one is still stripped here defensively so it never reaches an artefact.
 
+    Docstring blocks are located by :func:`_narrative_docstring_ranges`, the
+    single shared segmentation this module exposes, so a code string literal can
+    never be mistaken for one.
+
     This is the single shared strip layer: ``build_util.py_to_notebook`` routes
     both notebook generation (``generate.py``) and markdown generation
     (``generate_markdown.py``) through ``add_notebook_quotes``, and
-    ``navigator.py`` reuses this same tokenizer to segment docstrings — so
-    stripping here drops the section from every generated artefact and keeps it
-    out of the catalogue.
+    ``navigator.py`` reuses this same segmentation — it calls
+    ``add_notebook_quotes`` and reads the delimiters back out — so stripping here
+    drops the section from every generated artefact and keeps it out of the
+    catalogue.
     """
+    blocks = dict(_narrative_docstring_ranges(lines))
+
     out: List[str] = []
     i = 0
     n = len(lines)
@@ -46,17 +123,15 @@ def strip_env_declarations(lines: List[str]) -> List[str]:
             i += 1
             continue
 
-        # Docstring block: a bare `"""`/`'''` opener. Scan to its closing
-        # delimiter for a column-0 `__Env__` header appended anywhere inside.
-        if stripped in ('"""', "'''"):
-            delim = stripped
-            k = i + 1
+        # Docstring block: scan it for a column-0 `__Env__` header appended
+        # anywhere inside.
+        if i in blocks:
+            close = blocks[i]  # closing-delimiter index
             header = None
-            while k < n and lines[k].strip() != delim:
-                if header is None and lines[k].startswith("__Env__"):
+            for k in range(i + 1, close):
+                if lines[k].startswith("__Env__"):
                     header = k
-                k += 1
-            close = k  # closing-delimiter index (== n if unterminated)
+                    break
 
             if header is not None:
                 # Prose kept before the `__Env__` section, with the blank /
@@ -70,16 +145,15 @@ def strip_env_declarations(lines: List[str]) -> List[str]:
                     # preceding line.
                     out.append(line)
                     out.extend(kept)
-                    if close < n:
-                        out.append(lines[close])
+                    out.append(lines[close])
                 # else: the block holds only the `__Env__` section (standalone
                 # fallback) or is emptied by the strip — drop it whole.
-                i = close + 1 if close < n else n
+                i = close + 1
                 continue
 
             # A non-`__Env__` docstring block: emit it unchanged, delimiters too.
-            out.extend(lines[i : close + 1] if close < n else lines[i:])
-            i = close + 1 if close < n else n
+            out.extend(lines[i : close + 1])
+            i = close + 1
             continue
 
         out.append(line)
@@ -98,6 +172,12 @@ def add_notebook_quotes(lines: Iterable[str]):
     """
     Add %% above and below docs quotes with triple quotes.
 
+    Cell boundaries are the delimiter lines of the narrative docstring blocks
+    found by :func:`_narrative_docstring_ranges`, which parses the script rather
+    than testing line prefixes — so a triple-quoted string *assigned in code*,
+    whose closing delimiter also sits at column 0, is never mistaken for a
+    docstring boundary.
+
     A closing docstring does not emit its following code-cell marker until a
     non-blank code line is seen. This prevents adjacent docstrings from
     producing an empty code segment whose duplicate ``# %%`` markers are
@@ -143,13 +223,20 @@ def add_notebook_quotes(lines: Iterable[str]):
             f"authored. Delete them; the docstring blocks alone define the cells."
         )
 
+    # Re-derived on the stripped lines: dropping an `__Env__` block shifts every
+    # line number after it.
+    boundaries = set()
+    for start, end in _narrative_docstring_ranges(lines):
+        boundaries.add(start)
+        boundaries.add(end)
+
     out = list()
     is_in_quotes = False
     pending_code_boundary = False
     pending_lines: List[str] = []
 
-    for line in lines:
-        if line.startswith('"""') or line.startswith("'''"):
+    for index, line in enumerate(lines):
+        if index in boundaries:
             if is_in_quotes:
                 out.extend(["'''", "\n\n"])
                 pending_code_boundary = True

tests/test_add_notebook_quotes.py

@@ -170,3 +170,151 @@ def test_leading_docstring_does_not_produce_an_empty_first_code_cell(
     assert first["cell_type"] == "markdown"
     assert "__Intro__" in "".join(first["source"])
     assert [cell["cell_type"] for cell in notebook["cells"]].count("markdown") == 2
+
+
+# A triple-quoted string literal *assigned in code*. Its opener is indented
+# behind `s = ` and so is invisible to a line-prefix test, but its closer sits at
+# column 0 and does match one — flipping the docstring state and inverting every
+# cell boundary that follows. The code cell became an unterminated-string
+# SyntaxError and `print(s)` was emitted as narrative prose. Same failure class
+# as the opener bug (#211), opposite trigger.
+STRING_LITERAL_SCRIPT = (
+    '"""\n'
+    "__Intro__\n"
+    '"""\n'
+    "\n"
+    "x = 1\n"
+    's = """\n'
+    "literal\n"
+    '"""\n'
+    "print(s)\n"
+)
+
+
+# The same shape as it occurs in the wild: a gallery-build script whose
+# module-level `CSS` block is closed at column 0 and followed by real code.
+GALLERY_STYLE_SCRIPT = (
+    '"""\n'
+    "__Intro__\n"
+    '"""\n'
+    "\n"
+    'CSS = """\n'
+    "body { margin: 0; }\n"
+    '"""\n'
+    "\n"
+    "\n"
+    "def build():\n"
+    "    return CSS\n"
+)
+
+
+# The single-quote delimiter is equally affected; pin it so a fix that only
+# handles `\"\"\"` cannot pass.
+SINGLE_QUOTE_LITERAL_SCRIPT = (
+    '"""\n'
+    "__Intro__\n"
+    '"""\n'
+    "\n"
+    "s = '''\n"
+    "literal\n"
+    "'''\n"
+    "print(s)\n"
+)
+
+
+def test_code_string_literal_closer_is_not_a_cell_boundary():
+    converted = "".join(add_notebook_quotes(_lines(STRING_LITERAL_SCRIPT)))
+
+    # The literal must survive intact — delimiters unreplaced, no cell marker
+    # injected between them.
+    assert 's = """\nliteral\n"""\n' in converted
+    assert "print(s)" in converted
+
+    body = converted.split('s = """')[1]
+    assert "# %%" not in body.split('"""')[1]
+
+
+def test_code_string_literal_yields_markdown_then_intact_code_cell(
+    tmp_path, monkeypatch
+):
+    notebook = _notebook_from(
+        STRING_LITERAL_SCRIPT, tmp_path, monkeypatch, "string_literal.py"
+    )
+
+    assert [cell["cell_type"] for cell in notebook["cells"]] == ["markdown", "code"]
+
+    code = "".join(notebook["cells"][1]["source"])
+    assert 's = """\nliteral\n"""' in code
+    assert "print(s)" in code
+
+    for cell in notebook["cells"]:
+        if cell["cell_type"] == "code":
+            source = "".join(cell["source"])
+            assert "# %%" not in source
+            assert "'''" not in source
+
+
+def test_code_string_literal_code_cell_is_valid_python(tmp_path, monkeypatch):
+    """The regression's real symptom: the cell did not compile."""
+    import ast
+
+    notebook = _notebook_from(
+        STRING_LITERAL_SCRIPT, tmp_path, monkeypatch, "string_literal_parse.py"
+    )
+
+    for cell in notebook["cells"]:
+        if cell["cell_type"] == "code":
+            ast.parse("".join(cell["source"]))
+
+
+def test_gallery_style_css_block_is_not_a_cell_boundary(tmp_path, monkeypatch):
+    notebook = _notebook_from(
+        GALLERY_STYLE_SCRIPT, tmp_path, monkeypatch, "gallery_style.py"
+    )
+
+    assert [cell["cell_type"] for cell in notebook["cells"]] == ["markdown", "code"]
+
+    code = "".join(notebook["cells"][1]["source"])
+    assert 'CSS = """' in code
+    assert "def build():" in code
+
+
+def test_single_quoted_code_string_literal_is_not_a_cell_boundary():
+    converted = "".join(add_notebook_quotes(_lines(SINGLE_QUOTE_LITERAL_SCRIPT)))
+
+    assert "s = '''\nliteral\n'''\n" in converted
+    assert "print(s)" in converted
+
+
+def test_unparseable_source_raises():
+    import pytest
+
+    script = '"""\n__Intro__\n"""\n\ndef broken(:\n    pass\n'
+
+    with pytest.raises(ValueError, match="does not parse as Python"):
+        add_notebook_quotes(_lines(script))
+
+
+def test_single_line_docstring_raises():
+    import pytest
+
+    script = '"""__Intro__"""\n\nx = 1\n'
+
+    with pytest.raises(ValueError, match="single-line docstring"):
+        add_notebook_quotes(_lines(script))
+
+
+def test_indented_closing_delimiter_still_closes_the_block(tmp_path, monkeypatch):
+    """The mirror defect the prefix test also carried.
+
+    A closing delimiter written with leading whitespace did not match
+    ``startswith``, so the block never closed and every following line was
+    swallowed into the markdown cell. Deriving the span from the parsed source
+    closes it regardless of indentation.
+    """
+    script = '"""\n' "__Intro__\n" ' """\n' "\n" "x = 1\n"
+
+    notebook = _notebook_from(script, tmp_path, monkeypatch, "indented_closer.py")
+
+    assert [cell["cell_type"] for cell in notebook["cells"]] == ["markdown", "code"]
+    assert "x = 1" in "".join(notebook["cells"][1]["source"])

tests/test_strip_env_declarations.py

@@ -256,3 +256,47 @@ def test_navigator_ignores_env_section(tmp_path):
     title, summary, _ = navigator._parse_header(blocks[0])
     assert title == "Imaging Example"
     assert "one-line summary" in summary
+
+
+# An `__Env__` block followed by a triple-quoted string literal assigned in code.
+# The strip layer located docstring blocks by the same line-prefix test as the
+# notebook converter, so the literal's column-0 closer read as a block opener.
+# It was absorbed rather than acted on — every non-`__Env__` path emits the block
+# verbatim — but the misparse was real. Pin that the literal is untouched.
+ENV_THEN_STRING_LITERAL_SCRIPT = (
+    '"""\n'
+    "__Intro__\n"
+    "\n"
+    "__Env__\n"
+    "ENV: jax\n"
+    '"""\n'
+    "\n"
+    "x = 1\n"
+    's = """\n'
+    "literal\n"
+    '"""\n'
+    "print(s)\n"
+    "\n"
+    '"""\n'
+    "__Later__\n"
+    '"""\n'
+    "\n"
+    "y = 2\n"
+)
+
+
+def test_strip_leaves_a_code_string_literal_untouched():
+    stripped = "".join(
+        strip_env_declarations(_lines(ENV_THEN_STRING_LITERAL_SCRIPT))
+    )
+
+    # The `__Env__` section is gone, its docstring's prose kept.
+    assert "__Env__" not in stripped
+    assert "ENV: jax" not in stripped
+    assert "__Intro__" in stripped
+
+    # The literal and everything after it survive intact.
+    assert 's = """\nliteral\n"""\n' in stripped
+    assert "print(s)" in stripped
+    assert "__Later__" in stripped
+    assert "y = 2" in stripped