Process notebooks with pypandoc and a custom pandoc filter

Author: svekars

.github/workflows/build-tutorials.yml

@@ -187,7 +187,8 @@ jobs:
           echo "rm /opt/cache/bin/*" | docker exec -u root -i "${container_name}" bash
 
           docker exec -t "${container_name}" sh -c ".jenkins/build.sh"
-
+          docker exec -t "${container_name}" python "./post_process_notebooks.py"
+  
       - name: Upload docs preview
         uses: seemethere/upload-artifact-s3@v5
         if: ${{ github.event_name == 'pull_request' }}

conf.py

@@ -107,7 +107,10 @@ def reset_seeds(gallery_conf, fname):
                             "# https://pytorch.org/tutorials/beginner/colab\n"
                             "%matplotlib inline"),
     'reset_modules': (reset_seeds),
-    'ignore_pattern': r'_torch_export_nightly_tutorial.py'
+    'ignore_pattern': r'_torch_export_nightly_tutorial.py',
+        'pypandoc': {'extra_args': ['--mathjax'],
+                     'filters': ['./custom_pandoc_filter.py'],
+        },
 }
 
 if os.getenv('GALLERY_PATTERN'):

custom_pandoc_filter.py

@@ -0,0 +1,68 @@
+from pandocfilters import toJSONFilter, Div, RawBlock, Para, Str, Space, Link, Code, CodeBlock
+import markdown
+import re
+
+def to_markdown(item):
+    if item['t'] == 'Str':
+        return item['c']
+    elif item['t'] == 'Space':
+        return ' '
+    elif item['t'] == 'Link':
+        # Assuming the link text is always in the first item
+        return f"[{item['c'][1][0]['c']}]({item['c'][2][0]})"
+    elif item['t'] == 'Code':
+        return f"`{item['c'][1]}`"
+    elif item['t'] == 'CodeBlock':
+        return f"```\n{item['c'][1]}\n```"
+
+def process_admonitions(key, value, format, meta):
+    if key == 'Div':
+        [[ident, classes, keyvals], contents] = value
+        if 'note' in classes:
+            color = '#54c7ec'
+            label = 'NOTE:'
+        elif 'tip' in classes:
+            color = '#6bcebb'
+            label = 'TIP:'
+        elif 'warning' in classes:
+            color = '#e94f3b'
+            label = 'WARNING:'
+        else:
+            return
+
+        note_content = []
+        for block in contents:
+            if 't' in block and block['t'] == 'Para':
+                for item in block['c']:
+                    if item['t'] == 'Str':
+                        note_content.append(Str(item['c']))
+                    elif item['t'] == 'Space':
+                        note_content.append(Space())
+                    elif item['t'] == 'Link':
+                        note_content.append(Link(*item['c']))
+                    elif item['t'] == 'Code':
+                        note_content.append(Code(*item['c']))
+            elif 't' in block and block['t'] == 'CodeBlock':
+                note_content.append(CodeBlock(*block['c']))
+
+        note_content_md = ''.join(to_markdown(item) for item in note_content)
+        html_content = markdown.markdown(note_content_md)
+
+        return [{'t': 'RawBlock', 'c': ['html', f'<div style="background-color: {color}; color: #fff; font-weight: 700; padding-left: 10px; padding-top: 5px; padding-bottom: 5px">{label}</div>']}, {'t': 'RawBlock', 'c': ['html', '<div style="background-color: #f3f4f7; padding-left: 10px; padding-top: 10px; padding-bottom: 10px; padding-right: 10px">']}, {'t': 'RawBlock', 'c': ['html', html_content]}, {'t': 'RawBlock', 'c': ['html', '</div>']}]
+
+    elif key == 'RawBlock':
+        [format, content] = value
+        if format == 'html' and 'iframe' in content:
+            # Extract the video URL
+            video_url = content.split('src="')[1].split('"')[0]
+            # Create the Python code to display the video
+            html_code = f"""
+from IPython.display import display, HTML
+html_code = \"""
+{content}
+\"""
+display(HTML(html_code))
+"""
+
+if __name__ == "__main__":
+    toJSONFilter(process_admonitions)

post_process_notebooks.py

@@ -0,0 +1,46 @@
+import nbformat as nbf
+import os
+import re
+
+def get_gallery_dirs(conf_path):
+    """Execute the conf.py file and return the gallery directories."""
+    namespace = {}
+    exec(open(conf_path).read(), namespace)
+    sphinx_gallery_conf = namespace['sphinx_gallery_conf']
+    print(f"Processing directories: {', '.join(sphinx_gallery_conf['gallery_dirs'])}")
+    return sphinx_gallery_conf['gallery_dirs']
+
+def process_notebook(notebook_path):
+    """Read and process a notebook file."""
+    print(f'Processing file: {notebook_path}')
+    notebook = nbf.read(notebook_path, as_version=4)
+    for cell in notebook.cells:
+        if cell.cell_type == 'markdown':
+            cell.source = process_content(cell.source)
+    nbf.write(notebook, notebook_path)
+
+def process_content(content):
+    """Remove extra syntax from the content of a Markdown cell."""
+    content = re.sub(r'```{=html}\n<div', '<div', content)
+    content = re.sub(r'">\n```', '">', content)
+    content = re.sub(r'<\/div>\n```', '</div>\n', content)
+    content = re.sub(r'```{=html}\n</div>\n```', '</div>\n', content)
+    content = re.sub(r'```{=html}', '', content)
+    content = re.sub(r'</p>\n```', '</p>', content)
+    return content
+
+def process_directory(notebook_dir):
+    """Process all notebook files in a directory and its subdirectories."""
+    for root, dirs, files in os.walk(notebook_dir):
+        for filename in files:
+            if filename.endswith('.ipynb'):
+                process_notebook(os.path.join(root, filename))
+
+def main():
+    """Main function to process all directories specified in the conf.py file."""
+    conf_path = 'conf.py'
+    for notebook_dir in get_gallery_dirs(conf_path):
+        process_directory(notebook_dir)
+
+if __name__ == "__main__":
+    main()