Add a :formatter option to Kino.DataTable.new/2 (#441)

Co-authored-by: Jonatan Kłosko <[email protected]>

lib/kino/data_table.ex

@@ -43,14 +43,24 @@ defmodule Kino.DataTable do
       data. Sorting requires traversal of the whole enumerable, so it
       may not be desirable for large lazy enumerables. Defaults to `true`
 
+   * `:formatter` - a 2-arity function that is used to format the data
+     in the table. The first parameter passed is the `key` (column name) and
+     the second is the value to be formatted. When formatting column headings
+     the key is the special value `:__header__`. The formatter function must
+     return either `{:ok, string}` or `:default`. When the return value is
+     `:default` the default data table formatting is applied.
+
   """
   @spec new(Table.Reader.t(), keyword()) :: t()
   def new(tabular, opts \\ []) do
     name = Keyword.get(opts, :name, "Data")
     sorting_enabled = Keyword.get(opts, :sorting_enabled, true)
+    formatter = Keyword.get(opts, :formatter)
     {data_rows, data_columns, count, inspected} = prepare_data(tabular, opts)
 
-    Kino.Table.new(__MODULE__, {data_rows, data_columns, count, name, sorting_enabled, inspected},
+    Kino.Table.new(
+      __MODULE__,
+      {data_rows, data_columns, count, name, sorting_enabled, inspected, formatter},
       export: fn state -> {"text", state.inspected} end
     )
   end
@@ -162,7 +172,7 @@ defmodule Kino.DataTable do
   end
 
   @impl true
-  def init({data_rows, data_columns, count, name, sorting_enabled, inspected}) do
+  def init({data_rows, data_columns, count, name, sorting_enabled, inspected, formatter}) do
     features = Kino.Utils.truthy_keys(pagination: true, sorting: sorting_enabled)
     info = %{name: name, features: features}
 
@@ -174,8 +184,12 @@ defmodule Kino.DataTable do
        total_rows: count,
        slicing_fun: slicing_fun,
        slicing_cache: slicing_cache,
-       columns: Enum.map(data_columns, fn key -> %{key: key, label: value_to_string(key)} end),
-       inspected: inspected
+       columns:
+         Enum.map(data_columns, fn key ->
+           %{key: key, label: value_to_string(:__header__, key, formatter)}
+         end),
+       inspected: inspected,
+       formatter: formatter
      }}
   end
 
@@ -256,7 +270,9 @@ defmodule Kino.DataTable do
 
     data =
       Enum.map(records, fn record ->
-        Enum.map(state.columns, &(Map.fetch!(record, &1.key) |> value_to_string()))
+        Enum.map(state.columns, fn column ->
+          value_to_string(column.key, Map.fetch!(record, column.key), state.formatter)
+        end)
       end)
 
     total_rows = count || state.total_rows
@@ -279,6 +295,17 @@ defmodule Kino.DataTable do
     end
   end
 
+  defp value_to_string(_key, value, nil) do
+    value_to_string(value)
+  end
+
+  defp value_to_string(key, value, formatter) do
+    case formatter.(key, value) do
+      {:ok, string} -> string
+      :default -> value_to_string(value)
+    end
+  end
+
   defp value_to_string(value) when is_atom(value), do: inspect(value)
 
   defp value_to_string(value) when is_list(value) do
@@ -318,7 +345,10 @@ defmodule Kino.DataTable do
          total_rows: count,
          slicing_fun: slicing_fun,
          slicing_cache: slicing_cache,
-         columns: Enum.map(data_columns, fn key -> %{key: key, label: value_to_string(key)} end),
+         columns:
+           Enum.map(data_columns, fn key ->
+             %{key: key, label: value_to_string(:__header__, key, state.formatter)}
+           end),
          inspected: inspected
      }}
   end

test/kino/data_table_test.exs

@@ -275,6 +275,28 @@ defmodule Kino.DataTableTest do
     })
   end
 
+  test "supports a formatter option" do
+    entries = %{x: 1..3, y: [1.1, 1.2, 1.3]}
+
+    formatter =
+      fn
+        :__header__, value -> {:ok, "h:#{value}"}
+        :x, value when is_integer(value) -> {:ok, "x:#{value}"}
+        _, _ -> :default
+      end
+
+    kino = Kino.DataTable.new(entries, keys: [:x, :y], formatter: formatter)
+    data = connect(kino)
+
+    assert %{
+             content: %{
+               columns: [%{key: "0", label: "h:x"}, %{key: "1", label: "h:y"}],
+               data: [["x:1", "1.1"], ["x:2", "1.2"], ["x:3", "1.3"]],
+               total_rows: 3
+             }
+           } = data
+  end
+
   test "supports data update" do
     entries = [
       %User{id: 1, name: "Sherlock Holmes"},