udpates

frontend/src/components/variables/variables-table.tsx

@@ -1,10 +1,17 @@
 /* Copyright 2023 Marimo. All rights reserved. */
-import React from 'react'
-import { TableHeader, TableRow, TableHead, TableBody, TableCell, Table } from '../ui/table';
-import { Variables } from '@/core/variables/types';
-import { CellId } from '@/core/model/ids';
-import { CellLink } from '@/editor/links/cell-link';
-import { cn } from '@/lib/utils';
+import React from "react";
+import {
+  TableHeader,
+  TableRow,
+  TableHead,
+  TableBody,
+  TableCell,
+  Table,
+} from "../ui/table";
+import { Variables } from "@/core/variables/types";
+import { CellId } from "@/core/model/ids";
+import { CellLink } from "@/editor/links/cell-link";
+import { cn } from "@/lib/utils";
 
 interface Props {
   className?: string;
@@ -24,52 +31,78 @@ export const VariableTable: React.FC<Props> = ({
   cellIds.forEach((id, index) => cellIdToIndex.set(id, index));
 
   const sortedVariables = Object.values(variables).sort((a, b) => {
-    const aIndex = cellIdToIndex.get(a.declaredBy);
-    const bIndex = cellIdToIndex.get(b.declaredBy);
+    const aIndex = cellIdToIndex.get(a.declaredBy[0]);
+    const bIndex = cellIdToIndex.get(b.declaredBy[0]);
     if (aIndex === undefined || bIndex === undefined) {
       return 0;
     }
     return aIndex - bIndex;
   });
 
   return (
-    <Table className={cn("w-full overflow-hidden", className)}>
+    <Table className={cn("w-full overflow-hidden text-sm", className)}>
       <TableHeader>
-        <TableRow className='whitespace-nowrap'>
+        <TableRow className="whitespace-nowrap">
           <TableHead>Name</TableHead>
+          <TableHead>Type</TableHead>
           <TableHead>Declared In</TableHead>
           <TableHead>Used By</TableHead>
+          <TableHead>Value</TableHead>
         </TableRow>
       </TableHeader>
       <TableBody>
         {sortedVariables.map((variable) => (
           <TableRow key={variable.name}>
-            <TableCell className="font-medium max-w-[200px] text-ellipsis overflow-hidden"
+            <TableCell
+              className="font-medium max-w-[200px] text-ellipsis overflow-hidden"
               title={variable.name}
-            >{variable.name}</TableCell>
+            >
+              {variable.name}
+            </TableCell>
+            <TableCell className="font-medium max-w-[200px] text-ellipsis overflow-hidden">
+              {variable.dataType}
+            </TableCell>
             <TableCell>
-              <CellLink cellId={variable.declaredBy} />
+              {variable.declaredBy.length === 1 ? (
+                <CellLink cellId={variable.declaredBy[0]} />
+              ) : (
+                <div className="text-destructive flex flex-row gap-2">
+                  {variable.declaredBy.slice(0, 3).map((cellId, idx) => (
+                    <span key={cellId}>
+                      <CellLink
+                        key={cellId}
+                        cellId={cellId}
+                        className="whitespace-nowrap text-destructive"
+                      />
+                      {idx < variable.declaredBy.length - 1 && ", "}
+                    </span>
+                  ))}
+                </div>
+              )}
             </TableCell>
-            <TableCell className="flex flex-row overflow-auto">
+            <TableCell className="flex flex-row overflow-auto gap-2 items-baseline">
               {variable.usedBy.slice(0, 3).map((cellId, idx) => (
-                <React.Fragment key={cellId}>
+                <span key={cellId}>
                   <CellLink
                     key={cellId}
                     cellId={cellId}
-                    className="ml-2 whitespace-nowrap"
+                    className="whitespace-nowrap"
                   />
                   {idx < variable.usedBy.length - 1 && ", "}
-                </React.Fragment>
+                </span>
               ))}
               {variable.usedBy.length > 3 && (
-                <div className="ml-2 whitespace-nowrap text-muted-foreground text-sm">
+                <div className="whitespace-nowrap text-muted-foreground text-xs">
                   +{variable.usedBy.length - 3} more
                 </div>
               )}
             </TableCell>
+            <TableCell className="font-medium max-w-[200px] text-ellipsis overflow-hidden">
+              {variable.value}
+            </TableCell>
           </TableRow>
         ))}
       </TableBody>
     </Table>
   );
-}
+};

frontend/src/core/variables/__tests__/state.test.ts

@@ -0,0 +1,130 @@
+/* Copyright 2023 Marimo. All rights reserved. */
+import { beforeEach, describe, expect, it } from "vitest";
+import { CellId } from "@/core/model/ids";
+import { exportedForTesting } from "../state";
+import { VariableName, Variables } from "../types";
+
+const { initialState, reducer, createActions } = exportedForTesting;
+
+const CellIds = {
+  a: "a" as CellId,
+  b: "b" as CellId,
+};
+
+const Names = {
+  x: "x" as VariableName,
+  y: "y" as VariableName,
+};
+
+describe("cell reducer", () => {
+  let state: Variables;
+
+  const actions = createActions((action) => {
+    state = reducer(state, action);
+  });
+
+  beforeEach(() => {
+    state = initialState();
+  });
+
+  it("should set variables", () => {
+    const variables: Variables = {
+      [Names.x]: {
+        name: Names.x,
+        declaredBy: [CellIds.a],
+        usedBy: [CellIds.b],
+      },
+    };
+    actions.setVariables(variables);
+    expect(state).toEqual(variables);
+  });
+
+  it("should add variables", () => {
+    const x = {
+      name: Names.x,
+      declaredBy: [CellIds.a],
+      usedBy: [CellIds.b],
+    };
+    const y = {
+      name: Names.y,
+      declaredBy: [CellIds.a],
+      usedBy: [CellIds.b],
+    };
+    actions.addVariables([x, y]);
+    expect(state).toEqual({
+      [Names.x]: x,
+      [Names.y]: y,
+    });
+  });
+
+  it("should set metadata", () => {
+    const variables: Variables = {
+      [Names.x]: {
+        name: Names.x,
+        declaredBy: [CellIds.a],
+        usedBy: [CellIds.b],
+      },
+    };
+    actions.setVariables(variables);
+    expect(state).toEqual(variables);
+
+    // add metadata
+    actions.setMetadata([
+      {
+        name: Names.x,
+        value: "1",
+        dataType: "number",
+      },
+    ]);
+    expect(state).toEqual({
+      [Names.x]: {
+        name: Names.x,
+        declaredBy: [CellIds.a],
+        usedBy: [CellIds.b],
+        value: "1",
+        dataType: "number",
+      },
+    });
+
+    // drop unknown metadata
+    actions.setMetadata([
+      {
+        name: Names.x,
+        value: "2",
+        dataType: "number",
+      },
+      {
+        name: Names.y,
+        value: "3",
+        dataType: "number",
+      },
+    ]);
+    expect(state).toEqual({
+      [Names.x]: {
+        name: Names.x,
+        declaredBy: [CellIds.a],
+        usedBy: [CellIds.b],
+        value: "2",
+        dataType: "number",
+      },
+    });
+
+    // can re-add variables
+    actions.addVariables([
+      {
+        name: Names.x,
+        declaredBy: [CellIds.a],
+        usedBy: [],
+      },
+    ]);
+    expect(state).toEqual({
+      [Names.x]: {
+        name: Names.x,
+        declaredBy: [CellIds.a],
+        usedBy: [],
+        value: "2",
+        dataType: "number",
+      },
+    });
+  });
+});

frontend/src/core/variables/state.ts

@@ -1,7 +1,7 @@
 /* Copyright 2023 Marimo. All rights reserved. */
 
 import { createReducer } from "@/utils/createReducer";
-import { Variables } from "./types";
+import { Variable, VariableName, Variables } from "./types";
 import { atom, useAtomValue, useSetAtom } from "jotai";
 import { useMemo } from "react";
 
@@ -13,8 +13,33 @@ const { reducer, createActions } = createReducer(initialState, {
   setVariables: (_state, variables: Variables) => {
     return variables;
   },
-  addVariables: (state, variables: Variables) => {
-    return { ...state, ...variables };
+  addVariables: (state, variables: Variable[]) => {
+    const newVariables = { ...state };
+    for (const variable of variables) {
+      newVariables[variable.name] = {
+        ...newVariables[variable.name],
+        ...variable,
+      };
+    }
+    return newVariables;
+  },
+  setMetadata: (
+    state,
+    metadata: Array<{ name: VariableName; value?: string; dataType?: string }>
+  ) => {
+    const newVariables = { ...state };
+    for (const { name, value, dataType } of metadata) {
+      if (!newVariables[name]) {
+        continue;
+      }
+
+      newVariables[name] = {
+        ...newVariables[name],
+        value,
+        dataType: dataType,
+      };
+    }
+    return newVariables;
   },
 });
 
@@ -37,3 +62,9 @@ export function useVariablesActions() {
     return actions;
   }, [setState]);
 }
+
+export const exportedForTesting = {
+  reducer,
+  createActions,
+  initialState,
+};

frontend/src/core/variables/types.ts

@@ -7,8 +7,16 @@ export type VariableName = TypedString<"VariableName">;
 
 export interface Variable {
   name: VariableName;
-  declaredBy: CellId;
+  declaredBy: CellId[];
   usedBy: CellId[];
+  /**
+   * String representation of the value.
+   */
+  value?: string;
+  /**
+   * Type of the value.
+   */
+  dataType?: string;
 }
 
 export type Variables = Record<VariableName, Variable>;

frontend/src/editor/links/cell-link.tsx

@@ -16,7 +16,10 @@ export const CellLink = (props: Props): JSX.Element => {
 
   return (
     <div
-      className={cn("inline-block cursor-pointer text-blue-500 hover:underline", className)}
+      className={cn(
+        "inline-block cursor-pointer text-blue-500 hover:underline",
+        className
+      )}
       onClick={(e) => {
         e.stopPropagation();
         e.preventDefault();

frontend/src/stories/variables.stories.tsx

@@ -10,27 +10,38 @@ const meta: Meta<typeof VariableTable> = {
 };
 
 const variables = {
-  "a": {
-    "name": "a",
-    "declaredBy": "1",
-    "usedBy": ["2"],
+  a: {
+    name: "a",
+    declaredBy: ["1"],
+    usedBy: ["2"],
+    dataType: "number",
+    value: '1',
   },
-  "b": {
-    "name": "b",
-    "declaredBy": "2",
-    "usedBy": ["3"],
+  b: {
+    name: "b",
+    declaredBy: ["2"],
+    usedBy: ["3"],
+    dataType: "dataframe",
+    value: '<dataframe>',
   },
-  "my_super_super_long_variable_name": {
-    "name": "my_super_super_long_variable_name",
-    "declaredBy": "3",
-    "usedBy": Array.from({ length: 15 }, (_, i) => `${i + 4}`),
+  my_super_super_long_variable_name: {
+    name: "my_super_super_long_variable_name",
+    declaredBy: ["3"],
+    usedBy: Array.from({ length: 15 }, (_, i) => `${i + 4}`),
   },
-  "c": {
-    "name": "c",
-    "declaredBy": "4",
-    "usedBy": Array.from({ length: 3 }, (_, i) => `${i + 4}`),
+  c: {
+    name: "c",
+    declaredBy: ["4"],
+    dataType: "number",
+    value: '1',
+    usedBy: Array.from({ length: 3 }, (_, i) => `${i + 4}`),
   },
-}
+  d: {
+    name: "has_error",
+    declaredBy: ["4", "5"],
+    usedBy: Array.from({ length: 3 }, (_, i) => `${i + 4}`),
+  },
+};
 
 export default meta;
 type Story = StoryObj<typeof VariableTable>;