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

Author: jonatanklosko

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

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