From 401ca0d142333b7338cbaf4b790d986b763c777c Mon Sep 17 00:00:00 2001
From: Daniel Rodriguez <daniel@danielfrg.com>
Date: Sun, 19 Mar 2023 13:52:34 -0500
Subject: [PATCH] Add highlight_extra_classes to outputs

---
 mkdocs_jupyter/nbconvert2.py                      | 6 ++++++
 mkdocs_jupyter/tests/mkdocs/material-with-nbs.yml | 3 +--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/mkdocs_jupyter/nbconvert2.py b/mkdocs_jupyter/nbconvert2.py
index 9fef912..0f5a97e 100644
--- a/mkdocs_jupyter/nbconvert2.py
+++ b/mkdocs_jupyter/nbconvert2.py
@@ -155,6 +155,12 @@ def custom_clean_html(element):
             },
         )
 
+    if highlight_extra_classes:
+        content = content.replace(
+            "jp-OutputArea-output",
+            f"jp-OutputArea-output {highlight_extra_classes}",
+        )
+        return content
     return content
 
 
diff --git a/mkdocs_jupyter/tests/mkdocs/material-with-nbs.yml b/mkdocs_jupyter/tests/mkdocs/material-with-nbs.yml
index 806a288..ea02f0b 100644
--- a/mkdocs_jupyter/tests/mkdocs/material-with-nbs.yml
+++ b/mkdocs_jupyter/tests/mkdocs/material-with-nbs.yml
@@ -10,8 +10,7 @@ nav:
 plugins:
   - mkdocs-jupyter:
       include_source: true
-      highlight_extra_classes: "custom-css-classes"
-      include_requirejs: true
+      highlight_extra_classes: not-prose
 
 markdown_extensions:
   - toc: