[docs] Fix useStorageState regressions

docs/src/components/header/ThemeModeToggle.tsx

@@ -4,28 +4,25 @@ import IconButton from '@mui/material/IconButton';
 import Tooltip from '@mui/material/Tooltip';
 import DarkModeOutlined from '@mui/icons-material/DarkModeOutlined';
 import LightModeOutlined from '@mui/icons-material/LightModeOutlined';
-import useMediaQuery from '@mui/material/useMediaQuery';
-import { useChangeTheme } from 'docs/src/modules/components/ThemeContext';
+import { useColorSchemeShim } from 'docs/src/modules/components/ThemeContext';
 
-function CssVarsModeToggle(props: { onChange: (checked: boolean) => void }) {
-  const [mounted, setMounted] = React.useState(false);
+function CssVarsModeToggle(props: { onChange: (newMode: string) => void }) {
   const { mode, systemMode, setMode } = useColorScheme();
-  React.useEffect(() => {
-    setMounted(true);
-  }, []);
   const calculatedMode = mode === 'system' ? systemMode : mode;
+
   return (
     <Tooltip title={calculatedMode === 'dark' ? 'Turn on the light' : 'Turn off the light'}>
       <IconButton
         color="primary"
         disableTouchRipple
         disabled={!calculatedMode}
         onClick={() => {
-          props.onChange(calculatedMode === 'light');
-          setMode(calculatedMode === 'dark' ? 'light' : 'dark');
+          const newMode = calculatedMode === 'dark' ? 'light' : 'dark';
+          props.onChange(newMode);
+          setMode(newMode);
         }}
       >
-        {!calculatedMode || !mounted
+        {!calculatedMode
           ? null
           : {
               light: <DarkModeOutlined fontSize="small" />,
@@ -37,55 +34,37 @@ function CssVarsModeToggle(props: { onChange: (checked: boolean) => void }) {
 }
 
 export default function ThemeModeToggle() {
-  const theme = useTheme();
-  const changeTheme = useChangeTheme();
-  const [mode, setMode] = React.useState<string | null>(null);
-  const prefersDarkMode = useMediaQuery('(prefers-color-scheme: dark)');
-
-  React.useEffect(() => {
-    let initialMode = 'system';
-    try {
-      initialMode = localStorage.getItem('mui-mode') || initialMode;
-    } catch (error) {
-      // do nothing
-    }
-    setMode(initialMode);
-  }, []);
-
-  const handleChangeThemeMode = (checked: boolean) => {
-    const paletteMode = checked ? 'dark' : 'light';
-    setMode(paletteMode);
+  // TODO replace with useColorScheme once all pages support css vars
+  const { mode, systemMode, setMode } = useColorSchemeShim();
+  const calculatedMode = mode === 'system' ? systemMode : mode;
 
-    try {
-      localStorage.setItem('mui-mode', paletteMode); // syncing with homepage, can be removed once all pages are migrated to CSS variables
-    } catch (error) {
-      // do nothing
-    }
-    changeTheme({ paletteMode });
-  };
+  const theme = useTheme();
 
+  // Server-side hydration
   if (mode === null) {
     return <IconButton color="primary" disableTouchRipple />;
   }
 
-  if (theme.vars) {
-    // Temporarily renders conditionally because `useColorScheme` could not be used in the pages that haven't migrated to CSS theme variables.
-    return <CssVarsModeToggle onChange={handleChangeThemeMode} />;
+  // TODO remove this code branch, all pages should be migrated to use CssVarsProvider
+  if (!theme.vars) {
+    return (
+      <Tooltip title={calculatedMode === 'dark' ? 'Turn on the light' : 'Turn off the light'}>
+        <IconButton
+          color="primary"
+          disableTouchRipple
+          onClick={() => {
+            setMode(calculatedMode === 'dark' ? 'light' : 'dark');
+          }}
+        >
+          {calculatedMode === 'dark' ? (
+            <LightModeOutlined fontSize="small" />
+          ) : (
+            <DarkModeOutlined fontSize="small" />
+          )}
+        </IconButton>
+      </Tooltip>
+    );
   }
 
-  const checked = mode === 'system' ? prefersDarkMode : mode === 'dark';
-
-  return (
-    <Tooltip title={checked ? 'Turn on the light' : 'Turn off the light'}>
-      <IconButton
-        color="primary"
-        disableTouchRipple
-        onClick={() => {
-          handleChangeThemeMode(!checked);
-        }}
-      >
-        {checked ? <LightModeOutlined fontSize="small" /> : <DarkModeOutlined fontSize="small" />}
-      </IconButton>
-    </Tooltip>
-  );
+  return <CssVarsModeToggle onChange={setMode} />;
 }

docs/src/layouts/AppHeader.tsx

@@ -39,7 +39,6 @@ interface AppHeaderProps {
 
 export default function AppHeader(props: AppHeaderProps) {
   const { gitHubRepository = 'https://github.com/mui' } = props;
-
   const t = useTranslate();
 
   return (

docs/src/modules/components/AppSettingsDrawer.js

@@ -1,7 +1,6 @@
 import * as React from 'react';
 import PropTypes from 'prop-types';
 import { styled, useTheme } from '@mui/material/styles';
-import useMediaQuery from '@mui/material/useMediaQuery';
 import Drawer from '@mui/material/Drawer';
 import Box from '@mui/material/Box';
 import Button from '@mui/material/Button';
@@ -16,7 +15,7 @@ import DarkModeOutlinedIcon from '@mui/icons-material/DarkModeOutlined';
 import SettingsBrightnessIcon from '@mui/icons-material/SettingsBrightness';
 import FormatTextdirectionLToRIcon from '@mui/icons-material/FormatTextdirectionLToR';
 import FormatTextdirectionRToLIcon from '@mui/icons-material/FormatTextdirectionRToL';
-import { useChangeTheme } from 'docs/src/modules/components/ThemeContext';
+import { useColorSchemeShim, useChangeTheme } from 'docs/src/modules/components/ThemeContext';
 import { useTranslate } from '@mui/docs/i18n';
 
 const Heading = styled(Typography)(({ theme }) => ({
@@ -37,55 +36,26 @@ const IconToggleButton = styled(ToggleButton)({
   },
 });
 
-function AppSettingsDrawer(props) {
+export default function AppSettingsDrawer(props) {
   const { onClose, open = false, ...other } = props;
   const t = useTranslate();
   const upperTheme = useTheme();
   const changeTheme = useChangeTheme();
-  const [mode, setMode] = React.useState(null);
-  const prefersDarkMode = useMediaQuery('(prefers-color-scheme: dark)');
-  const preferredMode = prefersDarkMode ? 'dark' : 'light';
 
-  React.useEffect(() => {
-    // syncing with homepage, can be removed once all pages are migrated to CSS variables
-    let initialMode = 'system';
-    try {
-      initialMode = localStorage.getItem('mui-mode') || initialMode;
-    } catch (error) {
-      // do nothing
-    }
-    setMode(initialMode);
-  }, [preferredMode]);
+  // TODO replace with useColorScheme once all pages support css vars
+  const { mode, setMode } = useColorSchemeShim();
 
   const handleChangeThemeMode = (event, paletteMode) => {
     if (paletteMode === null) {
       return;
     }
-
     setMode(paletteMode);
-
-    if (paletteMode === 'system') {
-      try {
-        localStorage.setItem('mui-mode', 'system'); // syncing with homepage, can be removed once all pages are migrated to CSS variables
-      } catch (error) {
-        // thrown when cookies are disabled.
-      }
-      changeTheme({ paletteMode: preferredMode });
-    } else {
-      try {
-        localStorage.setItem('mui-mode', paletteMode); // syncing with homepage, can be removed once all pages are migrated to CSS variables
-      } catch (error) {
-        // thrown when cookies are disabled.
-      }
-      changeTheme({ paletteMode });
-    }
   };
 
   const handleChangeDirection = (event, direction) => {
     if (direction === null) {
       direction = upperTheme.direction;
     }
-
     changeTheme({ direction });
   };
 
@@ -202,5 +172,3 @@ AppSettingsDrawer.propTypes = {
   onClose: PropTypes.func.isRequired,
   open: PropTypes.bool,
 };
-
-export default AppSettingsDrawer;

docs/src/modules/components/HighlightedCodeWithTabs.tsx

@@ -4,6 +4,7 @@ import { Tabs, TabsOwnProps } from '@mui/base/Tabs';
 import { TabsList as TabsListBase } from '@mui/base/TabsList';
 import { TabPanel as TabPanelBase } from '@mui/base/TabPanel';
 import { Tab as TabBase } from '@mui/base/Tab';
+import useLocalStorageState from '@mui/utils/useLocalStorageState';
 import HighlightedCode from './HighlightedCode';
 
 const TabList = styled(TabsListBase)(({ theme }) => ({
@@ -77,44 +78,25 @@ type TabsConfig = {
   language: string;
   tab: string;
 };
-export default function HighlightedCodeWithTabs({
-  tabs,
-  storageKey,
-}: {
-  tabs: Array<TabsConfig>;
-  storageKey?: string;
-} & Record<string, any>) {
+
+export default function HighlightedCodeWithTabs(
+  props: {
+    tabs: Array<TabsConfig>;
+    storageKey?: string;
+  } & Record<string, any>,
+) {
+  const { tabs, storageKey } = props;
   const availableTabs = React.useMemo(() => tabs.map(({ tab }) => tab), [tabs]);
-  const [activeTab, setActiveTab] = React.useState(availableTabs[0]);
+  const [activeTab, setActiveTab] = useLocalStorageState(storageKey ?? null, availableTabs[0]);
 
   const [mounted, setMounted] = React.useState(false);
 
   React.useEffect(() => {
-    try {
-      setActiveTab((prev) => {
-        if (storageKey === undefined) {
-          return prev;
-        }
-        const storedValues = localStorage.getItem(storageKey);
-
-        return storedValues && availableTabs.includes(storedValues) ? storedValues : prev;
-      });
-    } catch (error) {
-      // ignore error
-    }
     setMounted(true);
-  }, [availableTabs, storageKey]);
+  }, []);
 
   const handleChange: TabsOwnProps['onChange'] = (event, newValue) => {
     setActiveTab(newValue as string);
-    if (storageKey === undefined) {
-      return;
-    }
-    try {
-      localStorage.setItem(storageKey, newValue as string);
-    } catch (error) {
-      // ignore error
-    }
   };
 
   const ownerState = { mounted };

docs/src/modules/components/MaterialYouUsageDemo.tsx

@@ -1,5 +1,5 @@
 import * as React from 'react';
-import { useTheme as md2UseTheme, alpha } from '@mui/material/styles';
+import { alpha } from '@mui/material/styles';
 import ReplayRoundedIcon from '@mui/icons-material/ReplayRounded';
 import Box from '@mui/material/Box';
 import Divider from '@mui/material/Divider';
@@ -13,13 +13,10 @@ import Typography from '@mui/material/Typography';
 import {
   extendTheme,
   CssVarsProvider as MaterialYouCssVarsProvider,
-  useColorScheme,
 } from '@mui/material-next/styles';
 import BrandingProvider from 'docs/src/BrandingProvider';
 import HighlightedCode from 'docs/src/modules/components/HighlightedCode';
 
-type Mode = 'light' | 'dark' | 'system';
-
 const materialYouTheme = extendTheme();
 const shallowEqual = (item1: { [k: string]: any }, item2: { [k: string]: any }) => {
   let equal = true;
@@ -93,14 +90,6 @@ export const prependLinesSpace = (code: string, size: number = 2) => {
   return newCode.join('\n');
 };
 
-function ModeSwitcher({ md2Mode }: { md2Mode: Mode }) {
-  const { setMode } = useColorScheme();
-  React.useEffect(() => {
-    setMode(md2Mode);
-  }, [md2Mode, setMode]);
-  return null;
-}
-
 interface MaterialYouUsageDemoProps<ComponentProps> {
   /**
    * Name of the component to show in the code block.
@@ -195,7 +184,6 @@ export default function MaterialYouUsageDemo<T extends { [k: string]: any } = {}
     }
   });
 
-  const md2Theme = md2UseTheme();
   return (
     <Box
       sx={{
@@ -221,7 +209,6 @@ export default function MaterialYouUsageDemo<T extends { [k: string]: any } = {}
           }}
         >
           <MaterialYouCssVarsProvider theme={materialYouTheme}>
-            <ModeSwitcher md2Mode={md2Theme.palette.mode} />
             {renderDemo(demoProps)}
           </MaterialYouCssVarsProvider>
         </Box>