Restore whether the variable is a dataframe on load (#159)

livebook-dev · Apr 9, 2024 · 18c1416 · 18c1416
1 parent c6e2e10
commit 18c1416
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 11 deletions.
diff --git a/lib/kino_explorer/data_transform_cell.ex b/lib/kino_explorer/data_transform_cell.ex
@@ -123,7 +123,8 @@ defmodule KinoExplorer.DataTransformCell do
     root_fields = %{
       "data_frame" => attrs["data_frame"],
       "assign_to" => attrs["assign_to"],
-      "collect" => if(Map.has_key?(attrs, "collect"), do: attrs["collect"], else: true)
+      "collect" => Map.get(attrs, "collect", true),
+      "is_data_frame" => Map.get(attrs, "is_data_frame", false)
     }
 
     operations = attrs["operations"]
@@ -205,7 +206,7 @@ defmodule KinoExplorer.DataTransformCell do
 
         _ ->
           %{
-            root_fields: ctx.assigns.root_fields,
+            root_fields: update_is_data_frame(ctx.assigns.root_fields, ctx),
             operations: update_data_options(ctx.assigns.operations, ctx)
           }
       end
@@ -347,14 +348,21 @@ defmodule KinoExplorer.DataTransformCell do
     %{
       root_fields: %{
         "data_frame" => data_frame,
+        "is_data_frame" => Map.get(ctx.assigns.data_frame_variables, data_frame, false),
         "assign_to" => nil,
-        "lazy" => true,
         "collect" => false
       },
       operations: default_operations() |> update_data_options(ctx, data_frame)
     }
   end
 
+  defp update_is_data_frame(%{"data_frame" => data_frame} = root_fields, ctx) do
+    case ctx.assigns.data_frame_variables do
+      %{^data_frame => is_data_frame} -> %{root_fields | "is_data_frame" => is_data_frame}
+      %{} -> root_fields
+    end
+  end
+
   defp updates_for_grouped_fields(:summarise, field, value, idx, ctx) do
     current_summarise = get_in(ctx.assigns.operations, [Access.at(idx)])
     columns = if field == "query", do: [], else: current_summarise["columns"]
@@ -434,7 +442,6 @@ defmodule KinoExplorer.DataTransformCell do
     |> Map.put("operations", ctx.assigns.operations)
     |> Map.put("data_frame_alias", ctx.assigns.data_frame_alias)
     |> Map.put("missing_require", ctx.assigns.missing_require)
-    |> Map.put("is_data_frame", is_data_frame?(ctx))
   end
 
   @impl true
@@ -939,7 +946,6 @@ defmodule KinoExplorer.DataTransformCell do
     |> Map.put("operations", partial_operations)
     |> Map.put("data_frame_alias", Explorer.DataFrame)
     |> Map.put("missing_require", Explorer.DataFrame)
-    |> Map.put("is_data_frame", is_data_frame?(ctx))
   end
 
   defp maybe_update_datalist(%{"operation_type" => "filters"} = operation, df) do
@@ -967,11 +973,6 @@ defmodule KinoExplorer.DataTransformCell do
 
   defp implements?(protocol, value), do: protocol.impl_for(value) != nil
 
-  defp is_data_frame?(ctx) do
-    df = ctx.assigns.root_fields["data_frame"]
-    Map.get(ctx.assigns.data_frame_variables, df)
-  end
-
   defp collect_index([%{name: :group_by}, %{name: :summarise} | rest], size, idx) do
     collect_index(rest, size, idx + 2)
   end

diff --git a/test/kino_explorer/data_transform_cell_test.exs b/test/kino_explorer/data_transform_cell_test.exs
@@ -63,6 +63,7 @@ defmodule KinoExplorer.DataTransformCellTest do
   @base_attrs %{
     "assign_to" => nil,
     "data_frame" => "teams",
+    "is_data_frame" => true,
     "collect" => true,
     "data_frame_alias" => Explorer.DataFrame,
     "missing_require" => nil,
@@ -75,6 +76,12 @@ defmodule KinoExplorer.DataTransformCellTest do
     assert source == ""
   end
 
+  test "if the variable is unknown, relies on attributse for is-dataframe check" do
+    {_kino, source} = start_smart_cell!(DataTransformCell, @base_attrs)
+
+    assert source == "teams"
+  end
+
   test "finds valid data in binding and sends the data options to the client" do
     {kino, _source} = start_smart_cell!(DataTransformCell, %{})
 
@@ -108,7 +115,7 @@ defmodule KinoExplorer.DataTransformCellTest do
             "value" => nil
           }
         ],
-        root_fields: %{"assign_to" => nil, "data_frame" => "people"}
+        root_fields: %{"assign_to" => nil, "data_frame" => "people", "is_data_frame" => true}
       }
     })
   end