feat(code-block): add word wrap functionality and controls to code bl… (#39689)

* feat(code-block): add word wrap functionality and controls to code blocks

* feat(code-block): enhance layout of code block with improved line number display and structure

* fix h-full issue

* fix Prettier issue

---------

Co-authored-by: Chris Chinchilla <chris.ward@supabase.io>

apps/docs/features/ui/CodeBlock/CodeBlock.client.tsx

@@ -1,10 +1,11 @@
 'use client'
 
-import { Check, Copy } from 'lucide-react'
-import { type MouseEvent, useCallback, useEffect, useState } from 'react'
+import { Check, Copy, WrapText, ArrowRightFromLine } from 'lucide-react'
+import { type MouseEvent, useCallback, useEffect, useState, useRef } from 'react'
 import { type ThemedToken } from 'shiki'
 import { type NodeHover } from 'twoslash'
 import { cn, copyToClipboard, Tooltip, TooltipContent, TooltipTrigger } from 'ui'
+import { getFontStyle } from './CodeBlock.utils'
 
 export function AnnotatedSpan({
   token,
@@ -113,3 +114,46 @@ export function CodeCopyButton({ className, content }: { className?: string; con
     </button>
   )
 }
+
+export function CodeBlockControls({ content }: { content: string }) {
+  const [isWrapped, setIsWrapped] = useState(false)
+  const wrapperRef = useRef<HTMLDivElement>(null)
+
+  const toggleWrap = useCallback(() => {
+    setIsWrapped((prev) => {
+      const newValue = !prev
+      // Find the parent code block and toggle the wrap data attribute
+      const codeBlock = wrapperRef.current?.closest('.shiki')
+      if (codeBlock) {
+        if (newValue) {
+          codeBlock.setAttribute('data-wrapped', 'true')
+        } else {
+          codeBlock.removeAttribute('data-wrapped')
+        }
+      }
+      return newValue
+    })
+  }, [])
+
+  return (
+    <div ref={wrapperRef} className="hidden group-hover:flex absolute top-2 right-2 gap-1">
+      <Tooltip>
+        <TooltipTrigger asChild>
+          <button
+            onClick={toggleWrap}
+            className={cn('border rounded-md p-1', 'hover:bg-selection transition')}
+            aria-label={isWrapped ? 'Disable word wrap' : 'Enable word wrap'}
+          >
+            {isWrapped ? (
+              <ArrowRightFromLine size={14} className="text-lighter" />
+            ) : (
+              <WrapText size={14} className="text-lighter" />
+            )}
+          </button>
+        </TooltipTrigger>
+        <TooltipContent>{isWrapped ? 'Disable word wrap' : 'Enable word wrap'}</TooltipContent>
+      </Tooltip>
+      <CodeCopyButton content={content} />
+    </div>
+  )
+}

apps/docs/features/ui/CodeBlock/CodeBlock.tsx

@@ -1,9 +1,9 @@
-import { type PropsWithChildren } from 'react'
+import { Fragment, type PropsWithChildren } from 'react'
 import { bundledLanguages, createHighlighter, type BundledLanguage, type ThemedToken } from 'shiki'
 import { createTwoslasher, type ExtraFiles, type NodeHover } from 'twoslash'
 import { cn } from 'ui'
 
-import { AnnotatedSpan, CodeCopyButton } from './CodeBlock.client'
+import { AnnotatedSpan, CodeBlockControls } from './CodeBlock.client'
 import { getFontStyle } from './CodeBlock.utils'
 import theme from './supabase-2.json' with { type: 'json' }
 import denoTypes from './types/lib.deno.d.ts.include'
@@ -72,27 +72,42 @@ export async function CodeBlock({
       )}
     >
       <pre>
-        <code className={lineNumbers ? 'flex' : ''}>
-          {lineNumbers && (
-            <div className="flex-shrink-0 select-none text-right text-muted bg-control py-6 px-2">
-              {tokens.map((_, idx) => (
-                <div key={idx} className="w-full">
-                  {idx + 1}
-                </div>
+        <code className={lineNumbers ? 'grid grid-cols-[auto_1fr]' : ''}>
+          {lineNumbers ? (
+            <>
+              {tokens.map((line, idx) => (
+                <Fragment key={idx}>
+                  <div
+                    className={cn(
+                      'select-none text-right text-muted bg-control px-2 min-h-5 leading-5',
+                      idx === 0 && 'pt-6',
+                      idx === tokens.length - 1 && 'pb-6'
+                    )}
+                  >
+                    {idx + 1}
+                  </div>
+                  <div
+                    className={cn(
+                      'code-content min-h-5 leading-5 pl-6 pr-6',
+                      idx === 0 && 'pt-6',
+                      idx === tokens.length - 1 && 'pb-6'
+                    )}
+                  >
+                    <CodeLine tokens={line} twoslash={twoslashed?.get(idx)} />
+                  </div>
+                </Fragment>
+              ))}
+            </>
+          ) : (
+            <div className="code-content p-6">
+              {tokens.map((line, idx) => (
+                <CodeLine key={idx} tokens={line} twoslash={twoslashed?.get(idx)} />
               ))}
             </div>
           )}
-          <div className={cn('p-6 overflow-x-auto', lineNumbers ? 'flex-grow' : '')}>
-            {tokens.map((line, idx) => (
-              <CodeLine key={idx} tokens={line} twoslash={twoslashed?.get(idx)} />
-            ))}
-          </div>
         </code>
       </pre>
-      <CodeCopyButton
-        content={code.trim()}
-        className="hidden group-hover:block absolute top-2 right-2"
-      />
+      <CodeBlockControls content={code.trim()} />
     </div>
   )
 }
@@ -112,7 +127,7 @@ function CodeLine({
   })
 
   return (
-    <span className="block h-5">
+    <span className="block min-h-5 leading-5">
       {tokens.map((token) =>
         twoslash?.has(token.offset) ? (
           <AnnotatedSpan

apps/docs/styles/main.scss

@@ -390,3 +390,10 @@ th code {
   --code-token-property: #15593b;
   --code-highlight-color: #1c1c1c;
 }
+
+/* Word wrap styles for code blocks */
+.shiki[data-wrapped='true'] .code-content {
+  overflow-x: hidden !important;
+  white-space: pre-wrap !important;
+  word-break: break-word !important;
+}