Sync convert_ipynb_to_mdx.py with Ax (#2795)

Summary:
Pull Request resolved: #2795

This script started to drift from the implementation in Ax. We need to keep these in sync when making changes to either repo.

Reviewed By: saitcakmak

Differential Revision: D72236261

fbshipit-source-id: dc08c33137bcc66bfb375c91ef16e647696d245f

Author: CristianLara

scripts/convert_ipynb_to_mdx.py

@@ -12,7 +12,7 @@
 import subprocess
 import uuid
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Union
 
 import mdformat
 import nbformat
@@ -45,7 +45,7 @@
 ]
 
 
-def load_nb_metadata() -> Dict[str, Dict[str, str]]:
+def load_nb_metadata() -> dict[str, dict[str, str]]:
     """
     Load the metadata and list of notebooks that are to be converted to MDX.
 
@@ -83,7 +83,7 @@ def load_notebook(path: Path) -> NotebookNode:
     return nb
 
 
-def create_folders(path: Path) -> Tuple[str, Path]:
+def create_folders(path: Path) -> tuple[str, Path]:
     """
     Create asset folders for the tutorial.
 
@@ -109,7 +109,7 @@ def create_folders(path: Path) -> Tuple[str, Path]:
     return filename, assets_folder
 
 
-def create_frontmatter(path: Path, nb_metadata: Dict[str, Dict[str, str]]) -> str:
+def create_frontmatter(path: Path, nb_metadata: dict[str, dict[str, str]]) -> str:
     """
     Create frontmatter for the resulting MDX file.
 
@@ -154,7 +154,7 @@ def create_imports() -> str:
     return f"{imports}\n"
 
 
-def get_current_git_tag() -> Optional[str]:
+def get_current_git_tag() -> str | None:
     """
     Retrieve the current Git tag if the current commit is tagged.
 
@@ -175,7 +175,7 @@ def get_current_git_tag() -> Optional[str]:
 
 
 def create_buttons(
-    nb_metadata: Dict[str, Dict[str, str]],
+    nb_metadata: dict[str, dict[str, str]],
 ) -> str:
     """
     Create buttons that link to Colab and GitHub for the tutorial.
@@ -196,10 +196,55 @@ def create_buttons(
     return f'<LinkButtons\n  githubUrl="{github_url}"\n  colabUrl="{colab_url}"\n/>\n\n'
 
 
-def handle_images_found_in_markdown(
+def handle_image_attachments(
+    markdown: str,
+    attachments: dict[str, dict[str, str]],
+) -> str:
+    """
+    Image attachments are stored in the notebook cell's "attachments" field in base64
+    format with their associated mime_type and referenced in the markdown via
+    attachment name.
+
+    The pattern we search for in the Markdown is
+    `![alt_text](attachment:attachment_name title)` with three groups:
+
+    - group 1 = alt_text (optional)
+    - group 2 = attachment_name
+    - group 3 = title (optional)
+
+    To represent this in MD we replace the attachment reference with the base64 encoded
+    string as `![{alt_text}](data:{mime_type};base64,{img_as_base64})`
+
+    Args:
+        markdown (str): The markdown content containing image attachments.
+        attachments (Dict[str, Dict[str, str]]): A dictionary of attachments with their
+            corresponding MIME types and base64 encoded data.
+
+    Returns:
+        str: The markdown content with images converted to base64 format.
+    """
+    markdown_image_pattern = re.compile(
+        r"""!\[([^\]]*)\]\(attachment:(.*?)\s*?(\".*\")?\)"""
+    )
+    # go through searches in reverse order so that each replacement doesn't affect the
+    # start/end indices for the next replacements
+    searches = reversed(list(re.finditer(markdown_image_pattern, markdown)))
+    for search in searches:
+        alt_text, attachment_name, _ = search.groups()
+        mime_type, base64 = next(iter(attachments[attachment_name].items()))
+        start, end = search.span()
+        markdown = (
+            markdown[:start]
+            + generate_img_base64_md(base64, mime_type, alt_text)
+            + markdown[end:]
+        )
+    return markdown
+
+
+def handle_image_paths_found_in_markdown(
     markdown: str,
     new_img_dir: Path,
-    lib_dir: Path,
+    nb_path: Path,
 ) -> str:
     """
     Update image paths in the Markdown, and copy the image to the docs location.
@@ -210,6 +255,9 @@ def handle_images_found_in_markdown(
     - group 1 = path/to/image.png
     - group 2 = "title"
 
+    We explicitly exclude matching if the path starts with `attachment:` as this
+    indicates that the image is embedded as a base64 attachment not a file path.
+
     The first group (the path to the image from the original notebook) will be replaced
     with ``assets/img/{name}`` where the name is `image.png` from the example above. The
     original image will also be copied to the new location
@@ -219,12 +267,15 @@ def handle_images_found_in_markdown(
         markdown (str): Markdown where we look for Markdown flavored images.
         new_img_dir (Path): Path where images are copied to for display in the
             MDX file.
-        lib_dir (Path): The location for the Bean Machine repo.
+        lib_dir (Path): The location for the repo.
+        nb_path (Path): The location for the notebook.
 
     Returns:
         str: The original Markdown with new paths for images.
     """
-    markdown_image_pattern = re.compile(r"""!\[[^\]]*\]\((.*?)(?=\"|\))(\".*\")?\)""")
+    markdown_image_pattern = re.compile(
+        r"""!\[[^\]]*\]\((?!attachment:)(.*?)(?=\"|\))(\".*\")?\)"""
+    )
     searches = list(re.finditer(markdown_image_pattern, markdown))
 
     # Return the given Markdown if no images are found.
@@ -250,11 +301,11 @@ def handle_images_found_in_markdown(
 
         # Copy the original image to the new location.
         if old_path.exists():
+            # resolves if an absolute path is used
             old_img_path = old_path
         else:
-            # Here we assume the original image exists in the same directory as the
-            # notebook, which should be in the tutorials folder of the library.
-            old_img_path = (lib_dir / "tutorials" / old_path).resolve()
+            # fall back to path relative to the notebook
+            old_img_path = (nb_path.parent / old_path).resolve()
         new_img_path = str(new_img_dir / name)
         shutil.copy(str(old_img_path), new_img_path)
 
@@ -359,7 +410,7 @@ def get_source(cell: NotebookNode) -> str:
 def handle_markdown_cell(
     cell: NotebookNode,
     new_img_dir: Path,
-    lib_dir: Path,
+    nb_path: Path,
 ) -> str:
     """
     Handle the given Jupyter Markdown cell and convert it to MDX.
@@ -368,17 +419,17 @@ def handle_markdown_cell(
         cell (NotebookNode): Jupyter Markdown cell object.
         new_img_dir (Path): Path where images are copied to for display in the
             Markdown cell.
-        lib_dir (Path): The location for the Bean Machine library.
+        lib_dir (Path): The location for the library.
+        nb_path (Path): The location for the notebook.
 
     Returns:
         str: Transformed Markdown object suitable for inclusion in MDX.
     """
     markdown = get_source(cell)
 
-    # Update image paths in the Markdown and copy them to the Markdown tutorials folder.
-    # Skip - Our images are base64 encoded, so we don't need to copy them to the docs
-    # folder.
-    # markdown = handle_images_found_in_markdown(markdown, new_img_dir, lib_dir)
+    # Handle the different ways images are included in the Markdown.
+    markdown = handle_image_paths_found_in_markdown(markdown, new_img_dir, nb_path)
+    markdown = handle_image_attachments(markdown, cell.get("attachments", {}))
 
     markdown = sanitize_mdx(markdown)
     mdx = mdformat.text(markdown, options={"wrap": 88}, extensions={"myst"})
@@ -411,9 +462,29 @@ def handle_cell_input(cell: NotebookNode, language: str) -> str:
     return f"```{language}\n{cell_source}\n```\n\n"
 
 
+def generate_img_base64_md(
+    img_as_base64: int | str | NotebookNode,
+    mime_type: int | str | NotebookNode,
+    alt_text: str = "",
+) -> str:
+    """
+    Generate a markdown image tag from a base64 encoded image.
+
+    Args:
+        img_as_base64 (int | str | NotebookNode): The base64 encoded image data.
+        mime_type (int | str | NotebookNode): The MIME type of the image.
+        alt_text (str, optional): The alternative text for the image. Defaults to an
+            empty string.
+
+    Returns:
+        str: A markdown formatted image tag.
+    """
+    return f"![{alt_text}](data:{mime_type};base64,{img_as_base64})"
+
+
 def handle_image(
-    values: List[Dict[str, Union[int, str, NotebookNode]]],
-) -> List[Tuple[int, str]]:
+    values: list[dict[str, int | str | NotebookNode]],
+) -> list[tuple[int, str]]:
     """
     Convert embedded images to string MDX can consume.
 
@@ -431,13 +502,13 @@ def handle_image(
         index = value["index"]
         mime_type = value["mime_type"]
         img = value["data"]
-        output.append((index, f"![](data:image/{mime_type};base64,{img})\n\n"))
+        output.append((index, f"{generate_img_base64_md(img, mime_type)}\n\n"))
     return output
 
 
 def handle_markdown(
-    values: List[Dict[str, Union[int, str, NotebookNode]]],
-) -> List[Tuple[int, str]]:
+    values: list[dict[str, int | str | NotebookNode]],
+) -> list[tuple[int, str]]:
     """
     Convert and format Markdown for MDX.
 
@@ -460,8 +531,8 @@ def handle_markdown(
 
 
 def handle_pandas(
-    values: List[Dict[str, Union[int, str, NotebookNode]]],
-) -> List[Tuple[int, str]]:
+    values: list[dict[str, int | str | NotebookNode]],
+) -> list[tuple[int, str]]:
     """
     Handle how to display pandas DataFrames.
 
@@ -503,11 +574,14 @@ def handle_pandas(
 
 
 def handle_plain(
-    values: List[Dict[str, Union[int, str, NotebookNode]]],
-) -> List[Tuple[int, str]]:
+    values: list[dict[str, int | str | NotebookNode]],
+) -> list[tuple[int, str]]:
     """
     Handle how to plain cell output should be displayed in MDX.
 
+    Stdout streams are chunked during execution, we merge adjacent streams here into
+    single cell output blocks.
+
     Args:
         values (List[Dict[str, Union[int, str, NotebookNode]]]): Bokeh tagged cell
             outputs.
@@ -518,6 +592,28 @@ def handle_plain(
             the tuple is the MDX formatted string.
     """
     output = []
+    adjacent_outputs = []
+    previous_index = -1
+
+    def append_to_output() -> None:
+        if not adjacent_outputs:
+            return
+        adjacent_outputs_str = "\n".join(adjacent_outputs)
+        output.append(
+            (
+                previous_index,
+                "\n".join(
+                    [
+                        "<CellOutput>",
+                        "{",
+                        f"`{adjacent_outputs_str}`",
+                        "}",
+                        "</CellOutput>\n\n",
+                    ]
+                ),
+            ),
+        )
+
     for value in values:
         index = int(value["index"])
         data = str(value["data"])
@@ -527,16 +623,23 @@ def handle_plain(
             data = "\n".join([line for line in str(value["data"]).splitlines() if line])
             # Remove backticks to make the text MDX compatible.
             data = data.replace("`", "")
-            output.append(
-                (index, f"<CellOutput>\n{{\n  `{data}`\n}}\n</CellOutput>\n\n"),
-            )
+            if previous_index in [-1, index - 1]:
+                # store in cache until we reach nonconsecutive index
+                adjacent_outputs.append(data)
+            else:
+                # flush cache to output and start a new one
+                append_to_output()
+                adjacent_outputs = [data]
+            previous_index = index
+    # flush the remaining cache to output
+    append_to_output()
     return output
 
 
 def handle_plotly(
-    values: List[Dict[str, Union[int, str, NotebookNode]]],
+    values: list[dict[str, int | str | NotebookNode]],
     plot_data_folder: Path,
-) -> List[Tuple[int, str]]:
+) -> list[tuple[int, str]]:
     """
     Convert Plotly outputs to MDX.
 
@@ -567,8 +670,8 @@ def handle_plotly(
 
 
 def handle_tqdm(
-    values: List[Dict[str, Union[int, str, NotebookNode]]],
-) -> List[Tuple[int, str]]:
+    values: list[dict[str, int | str | NotebookNode]],
+) -> list[tuple[int, str]]:
     """
     Handle the output of tqdm.
 
@@ -590,9 +693,9 @@ def handle_tqdm(
     return [(index, f"<CellOutput>\n{{\n  `{md}`\n}}\n</CellOutput>\n\n")]
 
 
-CELL_OUTPUTS_TO_PROCESS = Dict[
+CELL_OUTPUTS_TO_PROCESS = dict[
     str,
-    List[Dict[str, Union[int, str, NotebookNode]]],
+    list[dict[str, Union[int, str, NotebookNode]]],
 ]
 
 
@@ -637,8 +740,8 @@ def aggregate_mdx(
 
 
 def prioritize_dtypes(
-    cell_outputs: List[NotebookNode],
-) -> Tuple[List[List[str]], List[bool]]:
+    cell_outputs: list[NotebookNode],
+) -> tuple[list[list[str]], list[bool]]:
     """
     Prioritize cell output data types.
 
@@ -679,7 +782,7 @@ def aggregate_images_and_plotly(
     prioritized_data_dtype: str,
     cell_output: NotebookNode,
     data: NotebookNode,
-    plotly_flags: List[bool],
+    plotly_flags: list[bool],
     cell_outputs_to_process: CELL_OUTPUTS_TO_PROCESS,
     i: int,
 ) -> None:
@@ -742,7 +845,7 @@ def aggregate_plain_output(
     cell_outputs_to_process["plain"].append({"index": i, "data": data})
 
 
-def aggregate_output_types(cell_outputs: List[NotebookNode]) -> CELL_OUTPUTS_TO_PROCESS:
+def aggregate_output_types(cell_outputs: list[NotebookNode]) -> CELL_OUTPUTS_TO_PROCESS:
     """
     Aggregate cell outputs into a dictionary for further processing.
 
@@ -880,7 +983,7 @@ def transform_notebook(path: Path, nb_metadata: object) -> str:
 
         # Handle a Markdown cell.
         if cell_type == "markdown":
-            mdx += handle_markdown_cell(cell, img_folder, LIB_DIR)
+            mdx += handle_markdown_cell(cell, img_folder, path)
 
         # Handle a code cell.
         if cell_type == "code":