From 86ea3947ca6fdf61f603de92f2151f30a0fbfed2 Mon Sep 17 00:00:00 2001
From: ozakione <29860391+OzakIOne@users.noreply.github.com>
Date: Thu, 28 Sep 2023 18:03:42 +0200
Subject: [PATCH] refactor: sync wrap states in localstorage

---
 .../src/theme/CodeBlock/Content/String.tsx    |  1 +
 .../src/hooks/useCodeWordWrap.ts              | 38 ++++++++++---------
 2 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/packages/docusaurus-theme-classic/src/theme/CodeBlock/Content/String.tsx b/packages/docusaurus-theme-classic/src/theme/CodeBlock/Content/String.tsx
index 6a051802e2ef..122b14489dd0 100644
--- a/packages/docusaurus-theme-classic/src/theme/CodeBlock/Content/String.tsx
+++ b/packages/docusaurus-theme-classic/src/theme/CodeBlock/Content/String.tsx
@@ -85,6 +85,7 @@ export default function CodeBlockString({
               ref={wordWrap.codeBlockRef}
               className={clsx(className, styles.codeBlock, 'thin-scrollbar')}>
               <code
+                style={wordWrap.codeStyle}
                 className={clsx(
                   styles.codeBlockLines,
                   showLineNumbers && styles.codeBlockLinesWithNumbering,
diff --git a/packages/docusaurus-theme-common/src/hooks/useCodeWordWrap.ts b/packages/docusaurus-theme-common/src/hooks/useCodeWordWrap.ts
index c723e2d5d862..44779ee5cf75 100644
--- a/packages/docusaurus-theme-common/src/hooks/useCodeWordWrap.ts
+++ b/packages/docusaurus-theme-common/src/hooks/useCodeWordWrap.ts
@@ -4,7 +4,7 @@
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  */
-import type {RefObject} from 'react';
+import type {CSSProperties, RefObject} from 'react';
 import {useState, useCallback, useEffect, useRef} from 'react';
 import {useStorageSlot} from '../index';
 import {useMutationObserver} from './useMutationObserver';
@@ -53,30 +53,34 @@ function useTabBecameVisibleCallback(
   );
 }
 
+function useCodeWrapState() {
+  const [value, storageSlot] = useStorageSlot('docusaurus.code.wordWrap');
+
+  const toggle = useCallback(() => {
+    const newValue = value === 'true' ? 'false' : 'true';
+    storageSlot.set(newValue);
+  }, [value, storageSlot]);
+
+  return [value === 'true', toggle] as const;
+}
+
 export function useCodeWordWrap(): {
   readonly codeBlockRef: RefObject<HTMLPreElement>;
   readonly isEnabled: boolean;
   readonly isCodeScrollable: boolean;
   readonly toggle: () => void;
+  readonly codeStyle: CSSProperties;
 } {
-  const [value, storageSlot] = useStorageSlot('docusaurus.code.wordWrap');
+  const [isEnabled, toggleWrap] = useCodeWrapState();
   const [isCodeScrollable, setIsCodeScrollable] = useState<boolean>(false);
   const codeBlockRef = useRef<HTMLPreElement>(null);
+  const codeStyle: CSSProperties = isEnabled
+    ? {whiteSpace: 'pre-wrap', overflowWrap: 'anywhere'}
+    : {};
 
   const toggle = useCallback(() => {
-    const codeElement = codeBlockRef.current!.querySelector('code')!;
-
-    if (value === 'true') {
-      codeElement.removeAttribute('style');
-    } else {
-      codeElement.style.whiteSpace = 'pre-wrap';
-      // When code wrap is enabled, we want to avoid a scrollbar in any case
-      // Ensure that very very long words/strings/tokens still wrap
-      codeElement.style.overflowWrap = 'anywhere';
-    }
-
-    storageSlot.set(value === 'true' ? 'false' : 'true');
-  }, [codeBlockRef, value, storageSlot]);
+    toggleWrap();
+  }, [toggleWrap]);
 
   const updateCodeIsScrollable = useCallback(() => {
     const {scrollWidth, clientWidth} = codeBlockRef.current!;
@@ -90,7 +94,7 @@ export function useCodeWordWrap(): {
 
   useEffect(() => {
     updateCodeIsScrollable();
-  }, [value, updateCodeIsScrollable]);
+  }, [isEnabled, updateCodeIsScrollable]);
 
   useEffect(() => {
     window.addEventListener('resize', updateCodeIsScrollable, {
@@ -102,5 +106,5 @@ export function useCodeWordWrap(): {
     };
   }, [updateCodeIsScrollable]);
 
-  return {codeBlockRef, isEnabled: value === 'true', isCodeScrollable, toggle};
+  return {codeBlockRef, isEnabled, isCodeScrollable, toggle, codeStyle};
 }