Add store_line_numbers option for myst_to_notebook

jupytext/myst.py

@@ -35,15 +35,17 @@ def myst_extensions(no_md=False):
     return [".md", ".myst", ".mystnb", ".mnb"]
 
 
-def matches_mystnb(text, ext=None, requires_meta=True, require_non_md=True, return_nb=False):
+def matches_mystnb(
+    text, ext=None, requires_meta=True, require_non_md=True, return_nb=False
+):
     """Attempt to distinguish a file as mystnb, only given its extension and content.
 
     :param ext: the extension of the file
     :param requires_meta: requires the file to contain top matter metadata
     :param require_non_md: whether to require that a non-markdown cell is present
     :param return_nb: if a match is found, return the notebook
     """
-    if ext and "." + ("."+ext).rsplit(".", 1)[1] in myst_extensions(no_md=True):
+    if ext and "." + ("." + ext).rsplit(".", 1)[1] in myst_extensions(no_md=True):
         return True
     if requires_meta and not text.startswith("---"):
         return False
@@ -135,14 +137,20 @@ def _fmt_md(text):
 
 
 def myst_to_notebook(
-    text, code_directive=CODE_DIRECTIVE, raw_directive=RAW_DIRECTIVE, ignore_bad_meta=False
+    text,
+    code_directive=CODE_DIRECTIVE,
+    raw_directive=RAW_DIRECTIVE,
+    ignore_bad_meta=False,
+    store_line_numbers=False,
 ):
     """Convert text written in the myst format to a notebook.
 
     :param text: the file text
     :param code_directive: the name of the directive to search for containing code cells
     :param raw_directive: the name of the directive to search for containing raw cells
     :param ignore_bad_meta: ignore metadata that cannot be parsed as JSON/YAML
+    :param store_line_numbers: add a `_source_lines` key to cell metadata,
+        mapping to the source text.
 
     NOTE: we assume here that all of these directives are at the top-level,
     i.e. not nested in other directives.
@@ -196,9 +204,15 @@ def myst_to_notebook(
                     "".join(lines.lines[current_line:token.position.line_start - 1])
                 )
                 if source:
+                    md_metadata = nbf.from_dict(md_metadata)
+                    if store_line_numbers:
+                        md_metadata["_source_lines"] = [
+                            current_line,
+                            token.position.line_start - 1,
+                        ]
                     notebook.cells.append(
                         nbf_version.new_markdown_cell(
-                            source=source, metadata=nbf.from_dict(md_metadata),
+                            source=source, metadata=md_metadata,
                         )
                     )
                 if token.content:
@@ -251,35 +265,48 @@ def myst_to_notebook(
                     "".join(lines.lines[current_line:token.position.line_start - 1])
                 )
                 if md_source:
+                    md_metadata = nbf.from_dict(md_metadata)
+                    if store_line_numbers:
+                        md_metadata["_source_lines"] = [
+                            current_line,
+                            token.position.line_start - 1,
+                        ]
                     notebook.cells.append(
                         nbf_version.new_markdown_cell(
-                            source=md_source, metadata=nbf.from_dict(md_metadata),
+                            source=md_source, metadata=md_metadata,
                         )
                     )
                 current_line = token.position.line_end
                 md_metadata = {}
 
+                cell_metadata = nbf.from_dict(options)
+                if store_line_numbers:
+                    cell_metadata["_source_lines"] = [
+                        token.position.line_start,
+                        token.position.line_end,
+                    ]
                 if item.node.language == code_directive:
                     notebook.cells.append(
                         nbf_version.new_code_cell(
-                            source="\n".join(body_lines),
-                            metadata=nbf.from_dict(options),
+                            source="\n".join(body_lines), metadata=cell_metadata,
                         )
                     )
                 if item.node.language == raw_directive:
                     notebook.cells.append(
                         nbf_version.new_raw_cell(
-                            source="\n".join(body_lines),
-                            metadata=nbf.from_dict(options),
+                            source="\n".join(body_lines), metadata=cell_metadata,
                         )
                     )
 
         # add the final markdown cell (if present)
         if lines.lines[current_line:]:
+            md_metadata = nbf.from_dict(md_metadata)
+            if store_line_numbers:
+                md_metadata["_source_lines"] = [current_line, len(lines.lines)]
             notebook.cells.append(
                 nbf_version.new_markdown_cell(
                     source=_fmt_md("".join(lines.lines[current_line:])),
-                    metadata=nbf.from_dict(md_metadata),
+                    metadata=md_metadata,
                 )
             )
 

tests/test_ipynb_to_myst.py

@@ -4,13 +4,16 @@
     import mock
 from textwrap import dedent
 import pytest
+from nbformat import from_dict
+
+from jupytext.compare import compare_notebooks
 from jupytext.formats import get_format_implementation, JupytextFormatError
 from jupytext.myst import (
     myst_to_notebook,
     CODE_DIRECTIVE,
     MystMetadataParsingError,
     matches_mystnb,
-    myst_extensions
+    myst_extensions,
 )
 from .utils import requires_myst
 
@@ -105,6 +108,60 @@ def test_matches_mystnb():
 
 
 def test_not_installed():
-    with mock.patch('jupytext.formats.JUPYTEXT_FORMATS', return_value=[]):
+    with mock.patch("jupytext.formats.JUPYTEXT_FORMATS", return_value=[]):
         with pytest.raises(JupytextFormatError):
             get_format_implementation(".myst")
+
+
+def test_store_line_numbers():
+    notebook = myst_to_notebook(
+        dedent(
+            """\
+            ---
+            a: 1
+            ---
+            abc
+            +++
+            def
+            ```{{{0}}}
+            ---
+            b: 2
+            ---
+            c = 3
+            ```
+            xyz
+            """
+        ).format(CODE_DIRECTIVE),
+        store_line_numbers=True,
+    )
+    expected = {
+        "nbformat": 4,
+        "nbformat_minor": 4,
+        "metadata": {"a": 1},
+        "cells": [
+            {
+                "cell_type": "markdown",
+                "source": "abc",
+                "metadata": {"_source_lines": [3, 4]},
+            },
+            {
+                "cell_type": "markdown",
+                "source": "def",
+                "metadata": {"_source_lines": [5, 6]},
+            },
+            {
+                "cell_type": "code",
+                "metadata": {"b": 2, "_source_lines": [7, 12]},
+                "execution_count": None,
+                "source": "c = 3",
+                "outputs": [],
+            },
+            {
+                "cell_type": "markdown",
+                "source": "xyz",
+                "metadata": {"_source_lines": [12, 13]},
+            },
+        ],
+    }
+    notebook.nbformat_minor = 4
+    compare_notebooks(notebook, from_dict(expected))