Allow Kino.JS.Live outputs to be exported

lib/kino/js.ex

@@ -202,7 +202,7 @@ defmodule Kino.JS do
 
   defstruct [:module, :ref, :export]
 
-  @opaque t :: %__MODULE__{module: module(), ref: Kino.Output.ref(), export: map()}
+  @opaque t :: %__MODULE__{module: module(), ref: Kino.Output.ref(), export: boolean()}
 
   defmacro __using__(opts) do
     quote location: :keep, bind_quoted: [opts: opts] do
@@ -435,19 +435,23 @@ defmodule Kino.JS do
 
   ## Options
 
-    * `:export_info_string` - used as the info string for the Markdown
-      code block where output data is persisted
-
-    * `:export_key` - in case the data is a map and only a specific part
-      should be exported
+    * `:export` - a function called to export the given kino to Markdown.
+      See the "Export" section below
 
   ## Export
 
   The output can optionally be exported in notebook source by specifying
-  `:export_info_string`. For example:
+  an `:export` function. The function receives `data` as an argument
+  and should return a tuple `{info_string, payload}`. `info_string`
+  is used to annotate the Markdown code block where the output is
+  persisted. `payload` is the value persisted in the code block. The
+  value is automatically serialized to JSON, unless it is already a
+  string.
+
+  For example:
 
       data = "graph TD;A-->B;"
-      Kino.JS.new(__MODULE__, data, export_info_string: "mermaid")
+      Kino.JS.new(__MODULE__, data, export: fn data -> {"mermaid", data} end)
 
   Would be rendered as the following Live Markdown:
 
@@ -457,12 +461,22 @@ defmodule Kino.JS do
   ```
   ````
 
-  Non-binary data is automatically serialized to JSON.
+  > #### Export function {: .info}
+  >
+  > You should prefer to use the `data` argument for computing the
+  > export payload. However, if it cannot be inferred from `data`,
+  > you should just reference the original value. Do not put additional
+  > fields in `data`, just to use it for export.
   """
   @spec new(module(), term(), keyword()) :: t()
   def new(module, data, opts \\ []) do
+    # TODO: remove the old export options in Kino v0.14.0
     export =
       if info_string = opts[:export_info_string] do
+        IO.warn(
+          "passing :export_info_string to Kino.JS.new/3 is deprecated. Specify an :export function instead"
+        )
+
         export_key = opts[:export_key]
 
         if export_key do
@@ -477,17 +491,19 @@ defmodule Kino.JS do
           end
         end
 
-        %{info_string: info_string, key: export_key}
+        fn data -> {info_string, data[export_key]} end
       end
 
+    export = export || opts[:export]
+
     ref = Kino.Output.random_ref()
 
-    Kino.JS.DataStore.store(ref, data)
+    Kino.JS.DataStore.store(ref, data, export)
 
     Kino.Bridge.reference_object(ref, self())
     Kino.Bridge.monitor_object(ref, Kino.JS.DataStore, {:remove, ref})
 
-    %__MODULE__{module: module, ref: ref, export: export}
+    %__MODULE__{module: module, ref: ref, export: export != nil}
   end
 
   @doc false

lib/kino/js/data_store.ex

@@ -24,9 +24,9 @@ defmodule Kino.JS.DataStore do
   @doc """
   Stores output data under the given ref.
   """
-  @spec store(Kino.Output.ref(), term()) :: :ok
-  def store(ref, data) do
-    GenServer.cast(@name, {:store, ref, data})
+  @spec store(Kino.Output.ref(), term(), function()) :: :ok
+  def store(ref, data, export) do
+    GenServer.cast(@name, {:store, ref, data, export})
   end
 
   @impl true
@@ -35,19 +35,41 @@ defmodule Kino.JS.DataStore do
   end
 
   @impl true
-  def handle_cast({:store, ref, data}, state) do
-    {:noreply, put_in(state.ref_with_data[ref], data)}
+  def handle_cast({:store, ref, data, export}, state) do
+    state = put_in(state.ref_with_data[ref], %{data: data, export: export, export_result: nil})
+    {:noreply, state}
   end
 
   @impl true
   def handle_info({:connect, pid, %{origin: _origin, ref: ref}}, state) do
-    with {:ok, data} <- Map.fetch(state.ref_with_data, ref) do
+    with {:ok, %{data: data}} <- Map.fetch(state.ref_with_data, ref) do
       Kino.Bridge.send(pid, {:connect_reply, data, %{ref: ref}})
     end
 
     {:noreply, state}
   end
 
+  def handle_info({:export, pid, %{ref: ref}}, state) do
+    case state.ref_with_data do
+      %{^ref => info} ->
+        {state, export_result} =
+          if info.export_result do
+            {state, info.export_result}
+          else
+            export_result = info.export.(info.data)
+            state = put_in(state.ref_with_data[ref].export_result, export_result)
+            {state, export_result}
+          end
+
+        Kino.Bridge.send(pid, {:export_reply, export_result, %{ref: ref}})
+
+        {:noreply, state}
+
+      _ ->
+        {:noreply, state}
+    end
+  end
+
   def handle_info({:remove, ref}, state) do
     {_, state} = pop_in(state.ref_with_data[ref])
     {:noreply, state}

lib/kino/js/live.ex

@@ -182,11 +182,16 @@ defmodule Kino.JS.Live do
       end
   '''
 
-  defstruct [:module, :pid, :ref]
+  defstruct [:module, :pid, :ref, :export]
 
   alias Kino.JS.Live.Context
 
-  @opaque t :: %__MODULE__{module: module(), pid: pid(), ref: Kino.Output.ref()}
+  @opaque t :: %__MODULE__{
+            module: module(),
+            pid: pid(),
+            ref: Kino.Output.ref(),
+            export: boolean()
+          }
 
   @type payload :: term() | {:binary, info :: term(), binary()}
 
@@ -322,16 +327,25 @@ defmodule Kino.JS.Live do
 
   The given `init_arg` is passed to the `init/2` callback when
   the underlying kino process is started.
+
+  ## Options
+
+    * `:export` - a function called to export the given kino to Markdown.
+      This works the same as `Kino.JS.new/3`, except the function
+      receives `t:Kino.JS.Live.Context.t/0` as an argument
+
   """
-  @spec new(module(), term()) :: t()
-  def new(module, init_arg) do
+  @spec new(module(), term(), keyword()) :: t()
+  def new(module, init_arg, opts \\ []) do
+    export = opts[:export]
+
     ref = Kino.Output.random_ref()
 
-    case Kino.start_child({Kino.JS.Live.Server, {module, init_arg, ref}}) do
+    case Kino.start_child({Kino.JS.Live.Server, {module, init_arg, ref, export}}) do
       {:ok, pid} ->
         subscription_manager = Kino.SubscriptionManager.cross_node_name()
         Kino.Bridge.monitor_object(pid, subscription_manager, {:clear_topic, ref})
-        %__MODULE__{module: module, pid: pid, ref: ref}
+        %__MODULE__{module: module, pid: pid, ref: ref, export: export != nil}
 
       {:error, reason} ->
         raise ArgumentError,
@@ -348,7 +362,7 @@ defmodule Kino.JS.Live do
         pid: kino.pid,
         assets: kino.module.__assets_info__()
       },
-      export: nil
+      export: kino.export
     }
   end
 

lib/kino/js/live/server.ex

@@ -41,9 +41,9 @@ defmodule Kino.JS.Live.Server do
   end
 
   @impl true
-  def init({module, init_arg, ref}) do
+  def init({module, init_arg, ref, export}) do
     {:ok, ctx, _opts} = call_init(module, init_arg, ref)
-    {:ok, %{module: module, ctx: ctx}}
+    {:ok, %{module: module, ctx: ctx, export: export}}
   end
 
   @impl true
@@ -61,6 +61,12 @@ defmodule Kino.JS.Live.Server do
   end
 
   @impl true
+  def handle_info({:export, pid, %{}}, state) do
+    export_result = state.export.(state.ctx)
+    Kino.Bridge.send(pid, {:export_reply, export_result, %{ref: state.ctx.__private__.ref}})
+    {:noreply, state}
+  end
+
   def handle_info(msg, state) do
     case call_handle_info(msg, state.module, state.ctx) do
       {:ok, ctx} -> {:noreply, %{state | ctx: ctx}}

lib/kino/mermaid.ex

@@ -29,6 +29,6 @@ defmodule Kino.Mermaid do
   """
   @spec new(binary()) :: t()
   def new(content) do
-    Kino.JS.new(__MODULE__, content, export_info_string: "mermaid")
+    Kino.JS.new(__MODULE__, content, export: fn content -> {"mermaid", content} end)
   end
 end

lib/kino/table.ex

@@ -63,10 +63,22 @@ defmodule Kino.Table do
   @doc """
   Creates a new tabular kino using the given module as data
   specification.
+
+  ## Options
+
+    * `:export` - a function called to export the given kino to Markdown.
+      This works the same as `Kino.JS.new/3`, except the function
+      receives the state as an argument
+
   """
-  @spec new(module(), term()) :: t()
-  def new(module, init_arg) do
-    Kino.JS.Live.new(__MODULE__, {module, init_arg})
+  @spec new(module(), term(), keyword()) :: t()
+  def new(module, init_arg, opts \\ []) do
+    export =
+      if export = opts[:export] do
+        fn ctx -> export.(ctx.assigns.state) end
+      end
+
+    Kino.JS.Live.new(__MODULE__, {module, init_arg}, export: export)
   end
 
   @impl true

test/kino/js/data_store_test.exs

@@ -6,7 +6,7 @@ defmodule Kino.JS.DataStoreTest do
   test "replies to connect messages with stored data" do
     ref = Kino.Output.random_ref()
     data = [1, 2, 3]
-    DataStore.store(ref, data)
+    DataStore.store(ref, data, nil)
 
     send(DataStore, {:connect, self(), %{origin: "client1", ref: ref}})
     assert_receive {:connect_reply, ^data, %{ref: ^ref}}
@@ -15,11 +15,21 @@ defmodule Kino.JS.DataStoreTest do
   test "{:remove, ref} removes data for the given ref" do
     ref = Kino.Output.random_ref()
     data = [1, 2, 3]
-    DataStore.store(ref, data)
+    DataStore.store(ref, data, nil)
 
     send(DataStore, {:remove, ref})
 
     send(DataStore, {:connect, self(), %{origin: "client1", ref: ref}})
     refute_receive {:connect_reply, _, %{ref: ^ref}}
   end
+
+  test "replies to export messages when configured to" do
+    ref = Kino.Output.random_ref()
+    data = [1, 2, 3]
+    export = fn data -> {"text", inspect(data)} end
+    DataStore.store(ref, data, export)
+
+    send(DataStore, {:export, self(), %{origin: "client1", ref: ref}})
+    assert_receive {:export_reply, {"text", "[1, 2, 3]"}, %{ref: ^ref}}
+  end
 end

test/kino/js/live_test.exs

@@ -68,4 +68,10 @@ defmodule Kino.JS.LiveTest do
     Process.exit(pid, :kill)
     assert_receive {:DOWN, ^ref, :process, ^pid, :killed}
   end
+
+  test "export" do
+    %{ref: ref} = kino = LiveCounter.new(0)
+    send(kino.pid, {:export, self(), %{ref: ref}})
+    assert_receive {:export_reply, {"text", 0}, %{ref: ^ref}}
+  end
 end

test/kino/js_test.exs

@@ -16,43 +16,20 @@ defmodule Kino.JSTest do
   end
 
   describe "new/3" do
-    test "raises an error when :export_key is specified but data is not a map" do
-      assert_raise ArgumentError,
-                   "expected data to be a map, because :export_key is specified, got: []",
-                   fn ->
-                     Kino.JS.new(Kino.TestModules.JSExternalAssets, [],
-                       export_info_string: "lang",
-                       export_key: :spec
-                     )
-                   end
-    end
-
-    test "raises an error when :export_key not in data" do
-      assert_raise ArgumentError,
-                   "got :export_key of :spec, but no such key found in data: %{width: 10}",
-                   fn ->
-                     Kino.JS.new(Kino.TestModules.JSExternalAssets, %{width: 10},
-                       export_info_string: "lang",
-                       export_key: :spec
-                     )
-                   end
-    end
-
-    test "builds export info when :export_info_string is specified" do
+    test "sets export to true when :export is specified" do
       kino =
         Kino.JS.new(Kino.TestModules.JSExternalAssets, %{spec: %{"width" => 10, "height" => 10}},
-          export_info_string: "vega-lite",
-          export_key: :spec
+          export: fn vl -> {"vega-lite", vl.spec} end
         )
 
-      assert %{export: %{info_string: "vega-lite", key: :spec}} = kino
+      assert %{export: true} = kino
     end
 
-    test "sets export info to nil when :export_info_string is not specified" do
+    test "sets export to false when :export is not specified" do
       kino =
         Kino.JS.new(Kino.TestModules.JSExternalAssets, %{spec: %{"width" => 10, "height" => 10}})
 
-      assert %{export: nil} = kino
+      assert %{export: false} = kino
     end
   end
 end

test/support/test_modules/live_counter.ex

@@ -3,7 +3,7 @@ defmodule Kino.TestModules.LiveCounter do
   use Kino.JS.Live
 
   def new(count) do
-    Kino.JS.Live.new(__MODULE__, count)
+    Kino.JS.Live.new(__MODULE__, count, export: fn ctx -> {"text", ctx.assigns.count} end)
   end
 
   def bump(kino, by \\ 1) do