Moved ThemeProvider updates into effect

#1669

Author: bmingles

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}>

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 */

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);

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 : (
         <>

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]
   );