From 4dcd6db1e7e7b9aa7ecab75ba0abce5a3ea967de Mon Sep 17 00:00:00 2001
From: Brian Ingles <brianingles@deephaven.io>
Date: Fri, 8 Dec 2023 17:40:25 -0600
Subject: [PATCH] Moved ThemeProvider updates into effect

#1669
---
 packages/chart/src/ChartThemeProvider.tsx     | 20 +++++--------------
 packages/chart/src/MockChartModel.ts          |  2 +-
 .../src/RandomAreaPlotAnimation.tsx           | 15 ++++++--------
 .../components/src/theme/ThemeProvider.tsx    | 15 +++++++-------
 .../iris-grid/src/IrisGridThemeProvider.tsx   | 12 +++--------
 5 files changed, 23 insertions(+), 41 deletions(-)

diff --git a/packages/chart/src/ChartThemeProvider.tsx b/packages/chart/src/ChartThemeProvider.tsx
index cd5272e0f5..c13dc179c9 100644
--- a/packages/chart/src/ChartThemeProvider.tsx
+++ b/packages/chart/src/ChartThemeProvider.tsx
@@ -1,4 +1,4 @@
-import { createContext, ReactNode, useEffect, useState } from 'react';
+import { createContext, ReactNode, useMemo } from 'react';
 import { useTheme } from '@deephaven/components';
 import defaultChartTheme, { ChartTheme } from './ChartTheme';
 
@@ -20,20 +20,10 @@ export function ChartThemeProvider({
 }: ChartThemeProviderProps): JSX.Element {
   const { activeThemes } = useTheme();
 
-  const [chartTheme, setChartTheme] = useState<ChartTheme | null>(null);
-
-  // The `ThemeProvider` that supplies `activeThemes` also provides the corresponding
-  // CSS theme variables to the DOM by dynamically rendering <style> tags whenever
-  // the `activeThemes` change. Painting the latest CSS variables to the DOM may
-  // not happen until after `ChartThemeProvider` is rendered, but they should be
-  // available by the time the effect runs. Therefore, it is important to derive
-  // the chart theme in an effect instead of deriving in a `useMemo` to ensure
-  // we have the latest CSS variables.
-  useEffect(() => {
-    if (activeThemes != null) {
-      setChartTheme(defaultChartTheme());
-    }
-  }, [activeThemes]);
+  const chartTheme = useMemo(
+    () => (activeThemes == null ? null : defaultChartTheme()),
+    [activeThemes]
+  );
 
   return (
     <ChartThemeContext.Provider value={chartTheme}>
diff --git a/packages/chart/src/MockChartModel.ts b/packages/chart/src/MockChartModel.ts
index e9347bef99..ccac27ac36 100644
--- a/packages/chart/src/MockChartModel.ts
+++ b/packages/chart/src/MockChartModel.ts
@@ -20,7 +20,7 @@ class MockChartModel extends ChartModel {
 
   static smoothing = 1.5;
 
-  static _theme: ChartTheme;
+  static _theme: ChartTheme | null;
 
   static get theme(): ChartTheme {
     /* eslint-disable no-underscore-dangle */
diff --git a/packages/components/src/RandomAreaPlotAnimation.tsx b/packages/components/src/RandomAreaPlotAnimation.tsx
index 85bd75ef3f..2f4491880f 100644
--- a/packages/components/src/RandomAreaPlotAnimation.tsx
+++ b/packages/components/src/RandomAreaPlotAnimation.tsx
@@ -1,7 +1,7 @@
 /* eslint-disable react-hooks/exhaustive-deps */
 /* eslint-disable react/display-name */
 
-import React, { useEffect, useState, useRef } from 'react';
+import React, { useEffect, useState, useRef, useMemo } from 'react';
 import debounce from 'lodash.debounce';
 import { assertNotNull } from '@deephaven/utils';
 import './RandomAreaPlotAnimation.scss';
@@ -42,14 +42,11 @@ function getRandomAreaPlotAnimationThemeColors(): ThemeColors {
 const RandomAreaPlotAnimation = React.memo(() => {
   const { activeThemes } = useTheme();
 
-  const [themeColors, setThemeColors] = useState<ThemeColors | null>(null);
-
-  // Resolving css variables has to run in `useEffect` or `useLayoutEffect` so
-  // that we know React has updated the DOM with any styles set by the
-  // ThemeProvider.
-  useEffect(() => {
-    setThemeColors(getRandomAreaPlotAnimationThemeColors());
-  }, [activeThemes]);
+  const themeColors = useMemo(
+    () =>
+      activeThemes == null ? null : getRandomAreaPlotAnimationThemeColors(),
+    [activeThemes]
+  );
 
   const canvas = useRef<HTMLCanvasElement>(null);
   const container = useRef<HTMLDivElement>(null);
diff --git a/packages/components/src/theme/ThemeProvider.tsx b/packages/components/src/theme/ThemeProvider.tsx
index 9de00822ac..cfae6214f4 100644
--- a/packages/components/src/theme/ThemeProvider.tsx
+++ b/packages/components/src/theme/ThemeProvider.tsx
@@ -36,9 +36,11 @@ export interface ThemeProviderProps {
 export function ThemeProvider({
   themes: customThemes,
   children,
-}: ThemeProviderProps): JSX.Element {
+}: ThemeProviderProps): JSX.Element | null {
   const baseThemes = useMemo(() => getDefaultBaseThemes(), []);
 
+  const [value, setValue] = useState<ThemeContextValue | null>(null);
+
   const [selectedThemeKey, setSelectedThemeKey] = useState<string>(
     () => getThemePreloadData()?.themeKey ?? DEFAULT_DARK_THEME_KEY
   );
@@ -88,17 +90,16 @@ export function ThemeProvider({
     [activeThemes, selectedThemeKey, customThemes]
   );
 
-  const value = useMemo(
-    () => ({
+  useEffect(() => {
+    setValue({
       activeThemes,
       selectedThemeKey,
       themes,
       setSelectedThemeKey,
-    }),
-    [activeThemes, selectedThemeKey, themes]
-  );
+    });
+  }, [activeThemes, selectedThemeKey, themes]);
 
-  return (
+  return value == null ? null : (
     <ThemeContext.Provider value={value}>
       {activeThemes == null ? null : (
         <>
diff --git a/packages/iris-grid/src/IrisGridThemeProvider.tsx b/packages/iris-grid/src/IrisGridThemeProvider.tsx
index 38c315e127..775fc947aa 100644
--- a/packages/iris-grid/src/IrisGridThemeProvider.tsx
+++ b/packages/iris-grid/src/IrisGridThemeProvider.tsx
@@ -1,5 +1,5 @@
 import { useTheme } from '@deephaven/components';
-import { createContext, ReactNode, useEffect, useState } from 'react';
+import { createContext, ReactNode, useMemo } from 'react';
 import { createDefaultIrisGridTheme, IrisGridThemeType } from './IrisGridTheme';
 
 export type IrisGridThemeContextValue = Partial<IrisGridThemeType>;
@@ -16,14 +16,8 @@ export function IrisGridThemeProvider({
 }: IrisGridThemeProviderProps): JSX.Element {
   const { activeThemes } = useTheme();
 
-  const [gridTheme, setGridTheme] = useState<IrisGridThemeContextValue>({});
-
-  useEffect(
-    function refreshIrisGridTheme() {
-      if (activeThemes != null) {
-        setGridTheme(createDefaultIrisGridTheme());
-      }
-    },
+  const gridTheme = useMemo(
+    () => (activeThemes == null ? null : createDefaultIrisGridTheme()),
     [activeThemes]
   );