From 0b49d1784479208fe590b9c4d02fb1a80ee43246 Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Sat, 21 Dec 2024 14:03:32 +0800
Subject: [PATCH] add chart dark mode and colour based on status

---
 .../editor/chrome/panels/logs-panel.tsx       |  4 ++-
 .../src/components/tracing/tracing-spec.ts    | 33 ++++++++++---------
 frontend/src/components/tracing/tracing.tsx   | 24 +++++++++++---
 3 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/frontend/src/components/editor/chrome/panels/logs-panel.tsx b/frontend/src/components/editor/chrome/panels/logs-panel.tsx
index 81df21fbe56..40be0b1bdfb 100644
--- a/frontend/src/components/editor/chrome/panels/logs-panel.tsx
+++ b/frontend/src/components/editor/chrome/panels/logs-panel.tsx
@@ -79,7 +79,9 @@ function formatLog(log: CellLog) {
 
   return (
     <>
-      <span className="flex-shrink-0 text-[var(--gray-10)]">[{timestamp}]</span>
+      <span className="flex-shrink-0 text-[var(--gray-10)] dark:text-[var(--gray-11)]">
+        [{timestamp}]
+      </span>
       <span className={cn("flex-shrink-0", color)}>{level}</span>
       <span className="flex-shrink-0 text-[var(--gray-10)]">
         (<CellLink cellId={log.cellId} />)
diff --git a/frontend/src/components/tracing/tracing-spec.ts b/frontend/src/components/tracing/tracing-spec.ts
index 41cdc35cfb8..0233a007e0a 100644
--- a/frontend/src/components/tracing/tracing-spec.ts
+++ b/frontend/src/components/tracing/tracing-spec.ts
@@ -1,8 +1,8 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import type { CellId } from "@/core/cells/ids";
-import type { Field } from "@/plugins/impl/vega/types";
+import type { CellRun } from "@/core/cells/runs";
+import type { ResolvedTheme } from "@/theme/useTheme";
 import type { TopLevelSpec } from "vega-lite";
-import type { PositionDef } from "vega-lite/build/src/channeldef";
 
 export const REACT_HOVERED_CELLID = "hoveredCellId";
 export const VEGA_HOVER_SIGNAL = "cellHover";
@@ -14,31 +14,21 @@ export interface ChartValues {
   startTimestamp: string;
   endTimestamp: string;
   elapsedTime: string;
+  status: CellRun["status"];
 }
 
 export function createGanttBaseSpec(
   chartValues: ChartValues[],
   hiddenInputElementId: string,
   chartPosition: ChartPosition,
+  theme: ResolvedTheme,
 ): TopLevelSpec {
-  const yAxis: PositionDef<Field> | Partial<PositionDef<Field>> = {
-    field: "cellNum",
-    scale: { paddingInner: 0.2 },
-    sort: { field: "cellNum" },
-    title: "cell",
-  };
-
-  const hideYAxisLine = chartPosition === "sideBySide";
-  if (hideYAxisLine) {
-    yAxis.axis = null;
-  }
-
   return {
     $schema: "https://vega.github.io/schema/vega-lite/v5.json",
+    background: theme === "dark" ? "black" : undefined,
     mark: {
       type: "bar",
       cornerRadius: 2,
-      fill: "#37BE5F", // same colour as chrome's network tab
     },
     params: [
       {
@@ -56,7 +46,13 @@ export function createGanttBaseSpec(
     ],
     height: { step: 23 },
     encoding: {
-      y: yAxis,
+      y: {
+        field: "cellNum",
+        scale: { paddingInner: 0.2 },
+        sort: { field: "cellNum" },
+        title: "cell",
+        axis: chartPosition === "sideBySide" ? null : undefined,
+      },
       x: {
         field: "startTimestamp",
         type: "temporal",
@@ -82,6 +78,11 @@ export function createGanttBaseSpec(
           expr: `${REACT_HOVERED_CELLID} == toString(datum.cell) ? 19.5 : 18`,
         },
       },
+      color: {
+        field: "status",
+        scale: { domain: ["success", "error"], range: ["#37BE5F", "red"] }, // green is the same colour as chrome's network tab
+        legend: null,
+      },
     },
     data: {
       values: chartValues,
diff --git a/frontend/src/components/tracing/tracing.tsx b/frontend/src/components/tracing/tracing.tsx
index 8158e997735..abcf78caf97 100644
--- a/frontend/src/components/tracing/tracing.tsx
+++ b/frontend/src/components/tracing/tracing.tsx
@@ -34,11 +34,13 @@ import {
 import { ClearButton } from "../buttons/clear-button";
 import { cn } from "@/utils/cn";
 import { PanelEmptyState } from "../editor/chrome/panels/empty-state";
+import { type ResolvedTheme, useTheme } from "@/theme/useTheme";
 
 export const Tracing: React.FC = () => {
   const { runIds: newestToOldestRunIds, runMap } = useAtomValue(runsAtom);
   const { clearRuns } = useRunsActions();
 
+  const { theme } = useTheme();
   const [chartPosition, setChartPosition] = useState<ChartPosition>("above");
 
   const toggleChartPosition = () => {
@@ -88,6 +90,7 @@ export const Tracing: React.FC = () => {
                 key={run.runId}
                 run={run}
                 chartPosition={chartPosition}
+                theme={theme}
               />
             );
           }
@@ -106,6 +109,7 @@ interface ChartProps {
   height: number;
   vegaSpec: VisualizationSpec;
   signalListeners: SignalListeners;
+  theme: ResolvedTheme;
 }
 
 const Chart: React.FC<ChartProps> = (props: ChartProps) => {
@@ -113,6 +117,7 @@ const Chart: React.FC<ChartProps> = (props: ChartProps) => {
     <div className={props.className}>
       <LazyVega
         spec={props.vegaSpec}
+        theme={props.theme === "dark" ? "dark" : undefined}
         width={props.width}
         height={props.height}
         signalListeners={props.signalListeners}
@@ -129,7 +134,8 @@ interface VegaHoverCellSignal {
 const TraceBlock: React.FC<{
   run: Run;
   chartPosition: ChartPosition;
-}> = ({ run, chartPosition }) => {
+  theme: ResolvedTheme;
+}> = ({ run, chartPosition, theme }) => {
   const [collapsed, setCollapsed] = useState(false);
   const [hoveredCellId, setHoveredCellId] = useState<CellId | null>();
 
@@ -176,12 +182,18 @@ const TraceBlock: React.FC<{
       startTimestamp: formatChartTime(cellRun.startTime),
       endTimestamp: formatChartTime(cellRun.startTime + elapsedTime),
       elapsedTime: formatElapsedTime(elapsedTime * 1000),
+      status: cellRun.status,
     };
   });
 
   const hiddenInputElementId = `hiddenInputElement-${run.runId}`;
   const vegaSpec = compile(
-    createGanttBaseSpec(chartValues, hiddenInputElementId, chartPosition),
+    createGanttBaseSpec(
+      chartValues,
+      hiddenInputElementId,
+      chartPosition,
+      theme,
+    ),
   ).spec;
 
   const TraceRows = (
@@ -215,6 +227,7 @@ const TraceBlock: React.FC<{
               width={320}
               height={120}
               signalListeners={handleVegaSignal}
+              theme={theme}
             />
           )}
           {!collapsed && TraceRows}
@@ -236,6 +249,7 @@ const TraceBlock: React.FC<{
           width={240}
           height={100}
           signalListeners={handleVegaSignal}
+          theme={theme}
         />
       )}
     </div>
@@ -289,7 +303,7 @@ const TraceRow: React.FC<TraceRowProps> = ({
       onMouseEnter={handleMouseEnter}
       onMouseLeave={handleMouseLeave}
     >
-      <span className="text-[var(--gray-10)]">
+      <span className="text-[var(--gray-10)] dark:text-[var(--gray-11)]">
         [{formatLogTimestamp(cellRun.startTime)}]
       </span>
       <span className="text-[var(--gray-10)] w-16">
@@ -299,7 +313,9 @@ const TraceRow: React.FC<TraceRowProps> = ({
 
       <div className="flex flex-row gap-1 w-16 justify-end -ml-2">
         <Tooltip content={elapsedTimeTooltip}>
-          <span className="text-[var(--gray-10)]">{elapsedTimeStr}</span>
+          <span className="text-[var(--gray-10)] dark:text-[var(--gray-11)]">
+            {elapsedTimeStr}
+          </span>
         </Tooltip>
 
         <Tooltip content={cellRun.status}>