From 03f7a55aa00afa828f762503cff7a2b96d7b0c9b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Tue, 9 Apr 2024 19:09:13 +0800
Subject: [PATCH] Restore whether the variable is a dataframe on load

---
 lib/kino_explorer/data_transform_cell.ex      | 21 ++++++++++---------
 .../data_transform_cell_test.exs              |  9 +++++++-
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/lib/kino_explorer/data_transform_cell.ex b/lib/kino_explorer/data_transform_cell.ex
index 7d45858..5fdbd48 100644
--- 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
index d848814..f6e3509 100644
--- 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