Save static vegalite plot to livemd (#676)

* save static vegalite plot to livemd

* cleanup debug code

* using `vega-lite` as the type in the fenced code block

* wrap the text output in `{:text, output}` in take_outputs/2

* ignore :vega_lite_static when it is empty

* add import and export tests

* using `spec`

* format code

* keep the test focused

* improve tests for not including outputs

* always dump vage_lite spec

* Apply suggestions from code review

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>

lib/livebook/live_markdown/export.ex

@@ -153,6 +153,10 @@ defmodule Livebook.LiveMarkdown.Export do
     [delimiter, "output\n", text, "\n", delimiter]
   end
 
+  defp render_output({:vega_lite_static, spec}) do
+    ["```", "vega-lite\n", Jason.encode!(spec), "\n", "```"]
+  end
+
   defp render_output(_output), do: :ignored
 
   defp get_elixir_cell_code(%{source: source, disable_formatting: true}),

lib/livebook/live_markdown/import.ex

@@ -184,7 +184,17 @@ defmodule Livebook.LiveMarkdown.Import do
          [{"pre", _, [{"code", [{"class", "output"}], [output], %{}}], %{}} | ast],
          outputs
        ) do
-    take_outputs(ast, [output | outputs])
+    take_outputs(ast, [{:text, output} | outputs])
+  end
+
+  defp take_outputs(
+         [{"pre", _, [{"code", [{"class", "vega-lite"}], [output], %{}}], %{}} | ast],
+         outputs
+       ) do
+    case Jason.decode(output) do
+      {:ok, spec} -> take_outputs(ast, [{:vega_lite_static, spec} | outputs])
+      _ -> take_outputs(ast, outputs)
+    end
   end
 
   defp take_outputs(ast, outputs), do: {outputs, ast}
@@ -198,7 +208,6 @@ defmodule Livebook.LiveMarkdown.Import do
   defp build_notebook([{:cell, :elixir, source, outputs} | elems], cells, sections, messages) do
     {metadata, elems} = grab_metadata(elems)
     attrs = cell_metadata_to_attrs(:elixir, metadata)
-    outputs = Enum.map(outputs, &{:text, &1})
     cell = %{Notebook.Cell.new(:elixir) | source: source, outputs: outputs} |> Map.merge(attrs)
     build_notebook(elems, [cell | cells], sections, messages)
   end

test/livebook/live_markdown/export_test.exs

@@ -530,7 +530,13 @@ defmodule Livebook.LiveMarkdown.ExportTest do
                     | source: """
                       IO.puts("hey")\
                       """,
-                      outputs: ["hey"]
+                      outputs: [
+                        "hey",
+                        {:vega_lite_static,
+                         %{
+                           "$schema" => "https://vega.github.io/schema/vega-lite/v5.json"
+                         }}
+                      ]
                   }
                 ]
             }
@@ -650,7 +656,7 @@ defmodule Livebook.LiveMarkdown.ExportTest do
                     | source: """
                       IO.puts("hey")\
                       """,
-                      outputs: [{:vega_lite_static, %{}}, {:table_dynamic, self()}]
+                      outputs: [{:table_dynamic, self()}]
                   }
                 ]
             }
@@ -671,6 +677,75 @@ defmodule Livebook.LiveMarkdown.ExportTest do
 
       assert expected_document == document
     end
+
+    test "includes vega_lite_static output" do
+      notebook = %{
+        Notebook.new()
+        | name: "My Notebook",
+          sections: [
+            %{
+              Notebook.Section.new()
+              | name: "Section 1",
+                cells: [
+                  %{
+                    Notebook.Cell.new(:elixir)
+                    | source: """
+                      Vl.new(width: 500, height: 200)
+                      |> Vl.data_from_series(in: [1, 2, 3, 4, 5], out: [1, 2, 3, 4, 5])
+                      |> Vl.mark(:line)
+                      |> Vl.encode_field(:x, "in", type: :quantitative)
+                      |> Vl.encode_field(:y, "out", type: :quantitative)\
+                      """,
+                      outputs: [
+                        {:vega_lite_static,
+                         %{
+                           "$schema" => "https://vega.github.io/schema/vega-lite/v5.json",
+                           "data" => %{
+                             "values" => [
+                               %{"in" => 1, "out" => 1},
+                               %{"in" => 2, "out" => 2},
+                               %{"in" => 3, "out" => 3},
+                               %{"in" => 4, "out" => 4},
+                               %{"in" => 5, "out" => 5}
+                             ]
+                           },
+                           "encoding" => %{
+                             "x" => %{"field" => "in", "type" => "quantitative"},
+                             "y" => %{"field" => "out", "type" => "quantitative"}
+                           },
+                           "height" => 200,
+                           "mark" => "line",
+                           "width" => 500
+                         }}
+                      ]
+                  }
+                ]
+            }
+          ]
+      }
+
+      expected_document = """
+      # My Notebook
+
+      ## Section 1
+
+      ```elixir
+      Vl.new(width: 500, height: 200)
+      |> Vl.data_from_series(in: [1, 2, 3, 4, 5], out: [1, 2, 3, 4, 5])
+      |> Vl.mark(:line)
+      |> Vl.encode_field(:x, "in", type: :quantitative)
+      |> Vl.encode_field(:y, "out", type: :quantitative)
+      ```
+
+      ```vega-lite
+      {"$schema":"https://vega.github.io/schema/vega-lite/v5.json","data":{"values":[{"in":1,"out":1},{"in":2,"out":2},{"in":3,"out":3},{"in":4,"out":4},{"in":5,"out":5}]},"encoding":{"x":{"field":"in","type":"quantitative"},"y":{"field":"out","type":"quantitative"}},"height":200,"mark":"line","width":500}
+      ```
+      """
+
+      document = Export.notebook_to_markdown(notebook, include_outputs: true)
+
+      assert expected_document == document
+    end
   end
 
   test "includes outputs when notebook has :persist_outputs set" do

test/livebook/live_markdown/import_test.exs

@@ -582,6 +582,104 @@ defmodule Livebook.LiveMarkdown.ImportTest do
     assert %Notebook{name: "My Notebook", autosave_interval_s: 10} = notebook
   end
 
+  test "imports notebook with valid vega-lite output" do
+    markdown = """
+    # My Notebook
+
+    ## Section 1
+
+    ```elixir
+    Vl.new(width: 500, height: 200)
+    |> Vl.data_from_series(in: [1, 2, 3, 4, 5], out: [1, 2, 3, 4, 5])
+    |> Vl.mark(:line)
+    |> Vl.encode_field(:x, "in", type: :quantitative)
+    |> Vl.encode_field(:y, "out", type: :quantitative)
+    ```
+
+    ```vega-lite
+    {"$schema":"https://vega.github.io/schema/vega-lite/v5.json","data":{"values":[{"in":1,"out":1},{"in":2,"out":2},{"in":3,"out":3},{"in":4,"out":4},{"in":5,"out":5}]},"encoding":{"x":{"field":"in","type":"quantitative"},"y":{"field":"out","type":"quantitative"}},"height":200,"mark":"line","width":500}
+    ```
+    """
+
+    {notebook, []} = Import.notebook_from_markdown(markdown)
+
+    assert %Notebook{
+             name: "My Notebook",
+             sections: [
+               %Notebook.Section{
+                 name: "Section 1",
+                 cells: [
+                   %Cell.Elixir{
+                     source: """
+                     Vl.new(width: 500, height: 200)
+                     |> Vl.data_from_series(in: [1, 2, 3, 4, 5], out: [1, 2, 3, 4, 5])
+                     |> Vl.mark(:line)
+                     |> Vl.encode_field(:x, \"in\", type: :quantitative)
+                     |> Vl.encode_field(:y, \"out\", type: :quantitative)\
+                     """,
+                     outputs: [
+                       vega_lite_static: %{
+                         "$schema" => "https://vega.github.io/schema/vega-lite/v5.json",
+                         "data" => %{
+                           "values" => [
+                             %{"in" => 1, "out" => 1},
+                             %{"in" => 2, "out" => 2},
+                             %{"in" => 3, "out" => 3},
+                             %{"in" => 4, "out" => 4},
+                             %{"in" => 5, "out" => 5}
+                           ]
+                         },
+                         "encoding" => %{
+                           "x" => %{"field" => "in", "type" => "quantitative"},
+                           "y" => %{"field" => "out", "type" => "quantitative"}
+                         },
+                         "height" => 200,
+                         "mark" => "line",
+                         "width" => 500
+                       }
+                     ]
+                   }
+                 ]
+               }
+             ]
+           } = notebook
+  end
+
+  test "imports notebook with invalid vega-lite output" do
+    markdown = """
+    # My Notebook
+
+    ## Section 1
+
+    ```elixir
+    :ok
+    ```
+
+    ```vega-lite
+    not_a_json
+    ```
+    """
+
+    {notebook, []} = Import.notebook_from_markdown(markdown)
+
+    assert %Notebook{
+             name: "My Notebook",
+             sections: [
+               %Notebook.Section{
+                 name: "Section 1",
+                 cells: [
+                   %Cell.Elixir{
+                     source: """
+                     :ok\
+                     """,
+                     outputs: []
+                   }
+                 ]
+               }
+             ]
+           } = notebook
+  end
+
   test "skips invalid input type and returns a message" do
     markdown = """
     # My Notebook