Keep language class on livebook output code block (#1682)

elixir-lang · Mar 16, 2023 · 4b5c414 · 4b5c414
1 parent e63e957
commit 4b5c414
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/lib/ex_doc/markdown/earmark.ex b/lib/ex_doc/markdown/earmark.ex
@@ -81,9 +81,14 @@ defmodule ExDoc.Markdown.Earmark do
          ],
          acc
        ) do
-    code_attrs = Enum.reject(code_attrs, &match?({"class", _}, &1))
-    new_code = {"code", [{"class", "output"} | code_attrs], [source], code_meta}
-    fixup_list([{"pre", pre_attrs, [new_code], pre_meta} | ast], acc)
+    code_attrs =
+      case Enum.split_with(code_attrs, &match?({"class", _}, &1)) do
+        {[], attrs} -> [{"class", "output"} | attrs]
+        {[{"class", class}], attrs} -> [{"class", "#{class} output"} | attrs]
+      end
+
+    code_node = {"code", code_attrs, [source], code_meta}
+    fixup_list([{"pre", pre_attrs, [code_node], pre_meta} | ast], acc)
   end
 
   defp fixup_list([head | tail], acc) do

diff --git a/test/ex_doc/markdown/earmark_test.exs b/test/ex_doc/markdown/earmark_test.exs
@@ -57,13 +57,20 @@ defmodule ExDoc.Markdown.EarmarkTest do
       ```
       2
       ```
+
+      <!-- livebook:{"output":true} -->
+
+      ```mermaid
+      graph TD; A-->B;
+      ```
       """
 
       assert Markdown.to_ast(md, []) == [
                {:h1, [], ["Notebook"], %{}},
                {:h2, [], ["Example"], %{}},
                {:pre, [], [{:code, [class: "elixir"], ["1 + 1"], %{}}], %{}},
-               {:pre, [], [{:code, [class: "output"], ["2"], %{}}], %{}}
+               {:pre, [], [{:code, [class: "output"], ["2"], %{}}], %{}},
+               {:pre, [], [{:code, [class: "mermaid output"], ["graph TD; A-->B;"], %{}}], %{}}
              ]
     end
   end