Add VegaLite widget integration (#306)

* Add support for LiveWidget.VegaLite

* LiveWidget -> Kino

* Show an error when rendering unsupported Kino widget

* Match on Kino.Widget

* Add catch-all for unknown outputs

assets/js/vega_lite/index.js

@@ -1,13 +1,20 @@
+import * as vega from "vega";
 import vegaEmbed from "vega-embed";
 import { getAttributeOrThrow } from "../lib/attribute";
 
+// See https://github.com/vega/vega-lite/blob/b61b13c2cbd4ecde0448544aff6cdaea721fd22a/src/compile/data/assemble.ts#L228-L231
+const DEFAULT_DATASET_NAME = "source_0";
+
 /**
  * A hook used to render graphics according to the given
  * Vega-Lite specification.
  *
- * The hook expects a `vega_lite:<id>` event with `{ spec }` payload,
+ * The hook expects a `vega_lite:<id>:init` event with `{ spec }` payload,
  * where `spec` is the graphic definition as an object.
  *
+ * Later `vega_lite:<id>:push` events may be sent with `{ data, dataset, window }` payload,
+ * to dynamically update the underlying data.
+ *
  * Configuration:
  *
  *   * `data-id` - plot id
@@ -18,19 +25,45 @@ const VegaLite = {
     this.props = getProps(this);
     this.state = {
       container: null,
+      viewPromise: null,
     };
 
     this.state.container = document.createElement("div");
     this.el.appendChild(this.state.container);
 
-    this.handleEvent(`vega_lite:${this.props.id}`, ({ spec }) => {
+    this.handleEvent(`vega_lite:${this.props.id}:init`, ({ spec }) => {
-      vegaEmbed(this.state.container, spec, {});
+      if (!spec.data) {
+        spec.data = { values: [] };
+      }
+
+      this.state.viewPromise = vegaEmbed(this.state.container, spec, {}).then(
+        (result) => result.view
+      );
     });
+
+    this.handleEvent(
+      `vega_lite:${this.props.id}:push`,
+      ({ data, dataset, window }) => {
+        dataset = dataset || DEFAULT_DATASET_NAME;
+
+        this.state.viewPromise.then((view) => {
+          const currentData = view.data(dataset);
+          const changeset = buildChangeset(currentData, data, window);
+          view.change(dataset, changeset).run();
+        });
+      }
+    );
   },
 
   updated() {
     this.props = getProps(this);
   },
+
+  destroyed() {
+    if (this.state.viewPromise) {
+      this.state.viewPromise.then((view) => view.finalize());
+    }
+  },
 };
 
 function getProps(hook) {
@@ -39,4 +72,18 @@ function getProps(hook) {
   };
 }
 
+function buildChangeset(currentData, newData, window) {
+  if (window === 0) {
+    return vega.changeset().remove(currentData);
+  } else if (window) {
+    const toInsert = newData.slice(-window);
+    const freeSpace = Math.max(window - toInsert.length, 0);
+    const toRemove = currentData.slice(0, -freeSpace);
+
+    return vega.changeset().remove(toRemove).insert(toInsert);
+  } else {
+    return vega.changeset().insert(newData);
+  }
+}
+
 export default VegaLite;

lib/livebook/evaluator.ex

@@ -46,8 +46,8 @@ defmodule Livebook.Evaluator do
 
   Options:
 
-  * `formatter` - a module implementing the `Livebook.Evaluator.Formatter` behaviour,
+    * `formatter` - a module implementing the `Livebook.Evaluator.Formatter` behaviour,
-    used for transforming evaluation response before it's sent to the client
+      used for transforming evaluation response before it's sent to the client
   """
   def start_link(opts \\ []) do
     GenServer.start_link(__MODULE__, opts)
@@ -102,7 +102,7 @@ defmodule Livebook.Evaluator do
   def init(opts) do
     formatter = Keyword.get(opts, :formatter, Evaluator.IdentityFormatter)
 
-    {:ok, io_proxy} = Evaluator.IOProxy.start_link()
+    {:ok, io_proxy} = Evaluator.IOProxy.start_link(formatter: formatter)
 
     # Use the dedicated IO device as the group leader,
     # so that it handles all :stdio operations.
@@ -165,7 +165,7 @@ defmodule Livebook.Evaluator do
   end
 
   defp send_evaluation_response(send_to, ref, evaluation_response, formatter) do
-    response = formatter.format(evaluation_response)
+    response = formatter.format_response(evaluation_response)
     send(send_to, {:evaluation_response, ref, response})
   end
 

lib/livebook/evaluator/default_formatter.ex

@@ -6,36 +6,49 @@ defmodule Livebook.Evaluator.DefaultFormatter do
   @behaviour Livebook.Evaluator.Formatter
 
   @impl true
-  def format({:ok, :"do not show this result in output"}) do
+  def format_output(string) when is_binary(string), do: string
+  def format_output(other), do: format_value(other)
+
+  @impl true
+  def format_response({:ok, :"do not show this result in output"}) do
     # Functions in the `IEx.Helpers` module return this specific value
     # to indicate no result should be printed in the iex shell,
     # so we respect that as well.
     :ignored
   end
 
-  def format({:ok, {:module, _, _, _} = value}) do
+  def format_response({:ok, {:module, _, _, _} = value}) do
     inspected = inspect(value, inspect_opts(limit: 10))
     {:inspect, inspected}
   end
 
+  def format_response({:ok, value}) do
+    format_value(value)
+  end
+
+  def format_response({:error, kind, error, stacktrace}) do
+    formatted = Exception.format(kind, error, stacktrace)
+    {:error, formatted}
+  end
+
+  # ---
+
   @compile {:no_warn_undefined, {VegaLite, :to_spec, 1}}
 
-  def format({:ok, value}) do
+  defp format_value(value) do
     cond do
       is_struct(value, VegaLite) and function_exported?(VegaLite, :to_spec, 1) ->
         {:vega_lite_spec, VegaLite.to_spec(value)}
 
+      is_struct(value, Kino.Widget) ->
+        {:kino_widget, value.type, value.pid}
+
       true ->
         inspected = inspect(value, inspect_opts())
         {:inspect, inspected}
     end
   end
 
-  def format({:error, kind, error, stacktrace}) do
-    formatted = Exception.format(kind, error, stacktrace)
-    {:error, formatted}
-  end
-
   defp inspect_opts(opts \\ []) do
     default_opts = [pretty: true, width: 100, syntax_colors: syntax_colors()]
     Keyword.merge(default_opts, opts)

lib/livebook/evaluator/formatter.ex

@@ -9,11 +9,20 @@ defmodule Livebook.Evaluator.Formatter do
   # we would unnecessarily send a lot of data.
   # By defining a custom formatter the client can instruct
   # the `Evaluator` to send already transformed data.
+  #
+  # Additionally if the results rely on external package installed
+  # in the runtime node, then formatting anywhere else wouldn't be accurate,
+  # for example using `inspect` on an external struct.
 
   alias Livebook.Evaluator
 
+  @doc """
+  Transforms arbitrary evaluation output, usually binary.
+  """
+  @callback format_output(term()) :: term()
+
   @doc """
   Transforms the evaluation response.
   """
-  @callback format(Evaluator.evaluation_response()) :: term()
+  @callback format_response(Evaluator.evaluation_response()) :: term()
 end

lib/livebook/evaluator/identity_formatter.ex

@@ -1,10 +1,13 @@
 defmodule Livebook.Evaluator.IdentityFormatter do
   @moduledoc false
 
-  # The default formatter leaving the response unchanged.
+  # The default formatter leaving the output unchanged.
 
   @behaviour Livebook.Evaluator.Formatter
 
   @impl true
-  def format(evaluation_response), do: evaluation_response
+  def format_output(output), do: output
+
+  @impl true
+  def format_response(evaluation_response), do: evaluation_response
 end

lib/livebook/evaluator/io_proxy.ex

@@ -23,6 +23,11 @@ defmodule Livebook.Evaluator.IOProxy do
   Starts the IO device process.
 
   Make sure to use `configure/3` to actually proxy the requests.
+
+  Options:
+
+    * `formatter` - a module implementing the `Livebook.Evaluator.Formatter` behaviour,
+      used for transforming outputs before they are sent to the client
   """
   @spec start_link() :: GenServer.on_start()
   def start_link(opts \\ []) do
@@ -38,7 +43,7 @@ defmodule Livebook.Evaluator.IOProxy do
 
   The possible messages are:
 
-  * `{:evaluation_stdout, ref, string}` - for output requests,
+  * `{:evaluation_output, ref, string}` - for output requests,
     where `ref` is the given evaluation reference and `string` is the output.
   """
   @spec configure(pid(), pid(), Evaluator.ref()) :: :ok
@@ -57,8 +62,9 @@ defmodule Livebook.Evaluator.IOProxy do
   ## Callbacks
 
   @impl true
-  def init(_opts) do
+  def init(opts) do
-    {:ok, %{encoding: :unicode, target: nil, ref: nil, buffer: []}}
+    formatter = Keyword.get(opts, :formatter, Evaluator.IdentityFormatter)
+    {:ok, %{encoding: :unicode, target: nil, ref: nil, buffer: [], formatter: formatter}}
   end
 
   @impl true
@@ -150,6 +156,16 @@ defmodule Livebook.Evaluator.IOProxy do
     io_requests(reqs, {:ok, state})
   end
 
+  # Livebook custom request type, handled in a special manner
+  # by IOProxy and safely failing for any other IO device
+  # (resulting in the {:error, :request} response).
+  defp io_request({:livebook_put_term, term}, state) do
+    state = flush_buffer(state)
+    formatted_term = state.formatter.format_output(term)
+    send(state.target, {:evaluation_output, state.ref, formatted_term})
+    {:ok, state}
+  end
+
   defp io_request(_, state) do
     {{:error, :request}, state}
   end
@@ -186,7 +202,8 @@ defmodule Livebook.Evaluator.IOProxy do
     string = state.buffer |> Enum.reverse() |> Enum.join()
 
     if state.target != nil and string != "" do
-      send(state.target, {:evaluation_stdout, state.ref, string})
+      formatted_string = state.formatter.format_output(string)
+      send(state.target, {:evaluation_output, state.ref, formatted_string})
     end
 
     %{state | buffer: []}

lib/livebook/runtime.ex

@@ -52,8 +52,8 @@ defprotocol Livebook.Runtime do
   Evaluation outputs are send to the connected runtime owner.
   The messages should be of the form:
 
-  * `{:evaluation_stdout, ref, string}` - output captured during evaluation
+  * `{:evaluation_output, ref, output}` - output captured during evaluation
-  * `{:evaluation_response, ref, response}` - final result of the evaluation
+  * `{:evaluation_response, ref, output}` - final result of the evaluation
 
   If the evaluation state within a container is lost (e.g. a process goes down),
   the runtime can send `{:container_down, container_ref, message}` to notify the owner.

lib/livebook/session.ex

@@ -484,8 +484,8 @@ defmodule Livebook.Session do
     {:noreply, state}
   end
 
-  def handle_info({:evaluation_stdout, cell_id, string}, state) do
+  def handle_info({:evaluation_output, cell_id, string}, state) do
-    operation = {:add_cell_evaluation_stdout, self(), cell_id, string}
+    operation = {:add_cell_evaluation_output, self(), cell_id, string}
     {:noreply, handle_operation(state, operation)}
   end
 

lib/livebook/session/data.ex

@@ -27,7 +27,7 @@ defmodule Livebook.Session.Data do
     :users_map
   ]
 
-  alias Livebook.{Notebook, Evaluator, Delta, Runtime, JSInterop}
+  alias Livebook.{Notebook, Delta, Runtime, JSInterop}
   alias Livebook.Users.User
   alias Livebook.Notebook.{Cell, Section}
 
@@ -84,8 +84,8 @@ defmodule Livebook.Session.Data do
           | {:move_cell, pid(), Cell.id(), offset :: integer()}
           | {:move_section, pid(), Section.id(), offset :: integer()}
           | {:queue_cell_evaluation, pid(), Cell.id()}
-          | {:add_cell_evaluation_stdout, pid(), Cell.id(), String.t()}
+          | {:add_cell_evaluation_output, pid(), Cell.id(), term()}
-          | {:add_cell_evaluation_response, pid(), Cell.id(), Evaluator.evaluation_response()}
+          | {:add_cell_evaluation_response, pid(), Cell.id(), term()}
           | {:reflect_evaluation_failure, pid()}
           | {:cancel_cell_evaluation, pid(), Cell.id()}
           | {:set_notebook_name, pid(), String.t()}
@@ -265,23 +265,23 @@ defmodule Livebook.Session.Data do
     end
   end
 
-  def apply_operation(data, {:add_cell_evaluation_stdout, _client_pid, id, string}) do
+  def apply_operation(data, {:add_cell_evaluation_output, _client_pid, id, output}) do
     with {:ok, cell, _} <- Notebook.fetch_cell_and_section(data.notebook, id) do
       data
       |> with_actions()
-      |> add_cell_evaluation_stdout(cell, string)
+      |> add_cell_evaluation_output(cell, output)
       |> wrap_ok()
     else
       _ -> :error
     end
   end
 
-  def apply_operation(data, {:add_cell_evaluation_response, _client_pid, id, response}) do
+  def apply_operation(data, {:add_cell_evaluation_response, _client_pid, id, output}) do
     with {:ok, cell, section} <- Notebook.fetch_cell_and_section(data.notebook, id),
          :evaluating <- data.cell_infos[cell.id].evaluation_status do
       data
       |> with_actions()
-      |> add_cell_evaluation_response(cell, response)
+      |> add_cell_evaluation_response(cell, output)
       |> finish_cell_evaluation(cell, section)
       |> mark_dependent_cells_as_stale(cell)
       |> maybe_evaluate_queued()
@@ -532,22 +532,22 @@ defmodule Livebook.Session.Data do
     |> set_cell_info!(cell.id, evaluation_status: :ready)
   end
 
-  defp add_cell_evaluation_stdout({data, _} = data_actions, cell, string) do
+  defp add_cell_evaluation_output({data, _} = data_actions, cell, output) do
     data_actions
     |> set!(
       notebook:
         Notebook.update_cell(data.notebook, cell.id, fn cell ->
-          %{cell | outputs: add_output(cell.outputs, string)}
+          %{cell | outputs: add_output(cell.outputs, output)}
         end)
     )
   end
 
-  defp add_cell_evaluation_response({data, _} = data_actions, cell, response) do
+  defp add_cell_evaluation_response({data, _} = data_actions, cell, output) do
     data_actions
     |> set!(
       notebook:
         Notebook.update_cell(data.notebook, cell.id, fn cell ->
-          %{cell | outputs: add_output(cell.outputs, response)}
+          %{cell | outputs: add_output(cell.outputs, output)}
         end)
     )
   end

lib/livebook_web/live/kino/vega_lite_live.ex

@@ -0,0 +1,32 @@
+defmodule LivebookWeb.Kino.VegaLiteLive do
+  use LivebookWeb, :live_view
+
+  @impl true
+  def mount(_params, %{"pid" => pid, "id" => id}, socket) do
+    send(pid, {:connect, self()})
+
+    {:ok, assign(socket, id: id)}
+  end
+
+  @impl true
+  def render(assigns) do
+    ~L"""
+    <div id="vega-lite-<%= @id %>" phx-hook="VegaLite" phx-update="ignore" data-id="<%= @id %>">
+    </div>
+    """
+  end
+
+  @impl true
+  def handle_info({:connect_reply, %{spec: spec}}, socket) do
+    {:noreply, push_event(socket, "vega_lite:#{socket.assigns.id}:init", %{"spec" => spec})}
+  end
+
+  def handle_info({:push, %{data: data, dataset: dataset, window: window}}, socket) do
+    {:noreply,
+     push_event(socket, "vega_lite:#{socket.assigns.id}:push", %{
+       "data" => data,
+       "dataset" => dataset,
+       "window" => window
+     })}
+  end
+end

lib/livebook_web/live/session_live/cell_component.ex

@@ -228,7 +228,7 @@ defmodule LivebookWeb.SessionLive.CellComponent do
     <div class="flex flex-col rounded-lg border border-gray-200 divide-y divide-gray-200 font-editor">
       <%= for {output, index} <- @cell_view.outputs |> Enum.reverse() |> Enum.with_index(), output != :ignored do %>
         <div class="p-4 max-w-full overflow-y-auto tiny-scrollbar">
-          <%= render_output(socket, output, "#{@cell_view.id}-output#{index}") %>
+          <%= render_output(socket, output, "cell-#{@cell_view.id}-output#{index}") %>
         </div>
       <% end %>
     </div>
@@ -270,11 +270,35 @@ defmodule LivebookWeb.SessionLive.CellComponent do
     live_component(socket, LivebookWeb.SessionLive.VegaLiteComponent, id: id, spec: spec)
   end
 
+  defp render_output(socket, {:kino_widget, :vega_lite, pid}, id) do
+    live_render(socket, LivebookWeb.Kino.VegaLiteLive,
+      id: id,
+      session: %{"id" => id, "pid" => pid}
+    )
+  end
+
+  defp render_output(_socket, {:kino_widget, type, _pid}, _id) do
+    render_error_message_output("""
+    Got unsupported Kino widget type: #{inspect(type)}, if that's a new widget
+    make usre to update Livebook to the latest version
+    """)
+  end
+
   defp render_output(_socket, {:error, formatted}, _id) do
-    assigns = %{formatted: formatted}
+    render_error_message_output(formatted)
+  end
+
+  defp render_output(_socket, output, _id) do
+    # Above we cover all possible outputs from DefaultFormatter,
+    # but this is helpful in development when adding new output types.
+    render_error_message_output("Unknown output type: #{inspect(output)}")
+  end
+
+  defp render_error_message_output(message) do
+    assigns = %{message: message}
 
     ~L"""
-    <div class="overflow-auto whitespace-pre text-red-600 tiny-scrollbar"><%= @formatted %></div>
+    <div class="overflow-auto whitespace-pre text-red-600 tiny-scrollbar"><%= @message %></div>
     """
   end
 

lib/livebook_web/live/session_live/vega_lite_component.ex

@@ -9,13 +9,13 @@ defmodule LivebookWeb.SessionLive.VegaLiteComponent do
   @impl true
   def update(assigns, socket) do
     socket = assign(socket, id: assigns.id)
-    {:ok, push_event(socket, "vega_lite:#{socket.assigns.id}", %{"spec" => assigns.spec})}
+    {:ok, push_event(socket, "vega_lite:#{socket.assigns.id}:init", %{"spec" => assigns.spec})}
   end
 
   @impl true
   def render(assigns) do
     ~L"""
-    <div id="<%= @id %>" phx-hook="VegaLite" phx-update="ignore" data-id="<%= @id %>">
+    <div id="vega-lite-<%= @id %>" phx-hook="VegaLite" phx-update="ignore" data-id="<%= @id %>">
     </div>
     """
   end

test/livebook/evaluator/io_proxy_test.exs

@@ -12,17 +12,17 @@ defmodule Livebook.Evaluator.IOProxyTest do
   describe ":stdio interoperability" do
     test "IO.puts", %{io: io} do
       IO.puts(io, "hey")
-      assert_receive {:evaluation_stdout, :ref, "hey\n"}
+      assert_receive {:evaluation_output, :ref, "hey\n"}
     end
 
     test "IO.write", %{io: io} do
       IO.write(io, "hey")
-      assert_receive {:evaluation_stdout, :ref, "hey"}
+      assert_receive {:evaluation_output, :ref, "hey"}
     end
 
     test "IO.inspect", %{io: io} do
       IO.inspect(io, %{}, [])
-      assert_receive {:evaluation_stdout, :ref, "%{}\n"}
+      assert_receive {:evaluation_output, :ref, "%{}\n"}
     end
 
     test "IO.read", %{io: io} do
@@ -37,18 +37,26 @@ defmodule Livebook.Evaluator.IOProxyTest do
   test "buffers rapid output", %{io: io} do
     IO.puts(io, "hey")
     IO.puts(io, "hey")
-    assert_receive {:evaluation_stdout, :ref, "hey\nhey\n"}
+    assert_receive {:evaluation_output, :ref, "hey\nhey\n"}
   end
 
   test "respects CR as line cleaner", %{io: io} do
     IO.write(io, "hey")
     IO.write(io, "\roverride\r")
-    assert_receive {:evaluation_stdout, :ref, "\roverride\r"}
+    assert_receive {:evaluation_output, :ref, "\roverride\r"}
   end
 
   test "flush/1 synchronously sends buffer contents", %{io: io} do
     IO.puts(io, "hey")
     IOProxy.flush(io)
-    assert_received {:evaluation_stdout, :ref, "hey\n"}
+    assert_received {:evaluation_output, :ref, "hey\n"}
+  end
+
+  test "supports special livebook request type", %{io: io} do
+    ref = make_ref()
+    send(io, {:io_request, self(), ref, {:livebook_put_term, [1, 2, 3]}})
+    assert_receive {:io_reply, ^ref, :ok}
+
+    assert_received {:evaluation_output, :ref, [1, 2, 3]}
   end
 end

test/livebook/evaluator_test.exs

@@ -52,7 +52,7 @@ defmodule Livebook.EvaluatorTest do
     test "captures standard output and sends it to the caller", %{evaluator: evaluator} do
       Evaluator.evaluate_code(evaluator, self(), ~s{IO.puts("hey")}, :code_1)
 
-      assert_receive {:evaluation_stdout, :code_1, "hey\n"}
+      assert_receive {:evaluation_output, :code_1, "hey\n"}
     end
 
     test "using standard input results in an immediate error", %{evaluator: evaluator} do

test/livebook/runtime/erl_dist/manager_test.exs

@@ -63,7 +63,7 @@ defmodule Livebook.Runtime.ErlDist.ManagerTest do
 
       Manager.evaluate_code(node(), ~s{IO.puts(:stderr, "error")}, :container1, :evaluation1, nil)
 
-      assert_receive {:evaluation_stdout, :evaluation1, "error\n"}
+      assert_receive {:evaluation_output, :evaluation1, "error\n"}
 
       Manager.stop(node())
     end
@@ -80,7 +80,7 @@ defmodule Livebook.Runtime.ErlDist.ManagerTest do
 
       Manager.evaluate_code(node(), code, :container1, :evaluation1, nil)
 
-      assert_receive {:evaluation_stdout, :evaluation1, log_message}
+      assert_receive {:evaluation_output, :evaluation1, log_message}
       assert log_message =~ "[error] hey"
 
       Manager.stop(node())

test/livebook/session/data_test.exs

@@ -788,7 +788,7 @@ defmodule Livebook.Session.DataTest do
     end
   end
 
-  describe "apply_operation/2 given :add_cell_evaluation_stdout" do
+  describe "apply_operation/2 given :add_cell_evaluation_output" do
     test "updates the cell outputs" do
       data =
         data_after_operations!([
@@ -797,7 +797,7 @@ defmodule Livebook.Session.DataTest do
           {:queue_cell_evaluation, self(), "c1"}
         ])
 
-      operation = {:add_cell_evaluation_stdout, self(), "c1", "Hello!"}
+      operation = {:add_cell_evaluation_output, self(), "c1", "Hello!"}
 
       assert {:ok,
               %{
@@ -817,10 +817,10 @@ defmodule Livebook.Session.DataTest do
           {:insert_section, self(), 0, "s1"},
           {:insert_cell, self(), "s1", 0, :elixir, "c1"},
           {:queue_cell_evaluation, self(), "c1"},
-          {:add_cell_evaluation_stdout, self(), "c1", "Hola"}
+          {:add_cell_evaluation_output, self(), "c1", "Hola"}
         ])
 
-      operation = {:add_cell_evaluation_stdout, self(), "c1", " amigo!"}
+      operation = {:add_cell_evaluation_output, self(), "c1", " amigo!"}
 
       assert {:ok,
               %{
@@ -840,10 +840,10 @@ defmodule Livebook.Session.DataTest do
           {:insert_section, self(), 0, "s1"},
           {:insert_cell, self(), "s1", 0, :elixir, "c1"},
           {:queue_cell_evaluation, self(), "c1"},
-          {:add_cell_evaluation_stdout, self(), "c1", "Hola"}
+          {:add_cell_evaluation_output, self(), "c1", "Hola"}
         ])
 
-      operation = {:add_cell_evaluation_stdout, self(), "c1", "\ramigo!\r"}
+      operation = {:add_cell_evaluation_output, self(), "c1", "\ramigo!\r"}
 
       assert {:ok,
               %{
@@ -866,7 +866,7 @@ defmodule Livebook.Session.DataTest do
           {:add_cell_evaluation_response, self(), "c1", {:ok, [1, 2, 3]}}
         ])
 
-      operation = {:add_cell_evaluation_stdout, self(), "c1", "Hello!"}
+      operation = {:add_cell_evaluation_output, self(), "c1", "Hello!"}
 
       assert {:ok,
               %{