improvement: add copy button to codeblocks

frontend/src/components/icons/copy-icon.tsx

@@ -8,7 +8,7 @@ import { Events } from "@/utils/events";
 import { copyToClipboard } from "@/utils/copy";
 
 interface Props {
-  value: string;
+  value: string | (() => string);
   className?: string;
   tooltip?: string | false;
 }
@@ -21,7 +21,8 @@ export const CopyClipboardIcon: React.FC<Props> = ({
   const [isCopied, setIsCopied] = useState(false);
 
   const handleCopy = Events.stopPropagation(async () => {
-    await copyToClipboard(value).then(() => {
+    const valueToCopy = typeof value === "function" ? value() : value;
+    await copyToClipboard(valueToCopy).then(() => {
       setIsCopied(true);
       setTimeout(() => setIsCopied(false), 2000);
     });

frontend/src/plugins/core/RenderHTML.tsx

@@ -1,12 +1,18 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import parse, { Element, type DOMNode } from "html-react-parser";
-import React from "react";
+import { CopyClipboardIcon } from "@/components/icons/copy-icon";
+import parse, {
+  Element,
+  type HTMLReactParserOptions,
+  type DOMNode,
+} from "html-react-parser";
+import React, { useId, type ReactNode } from "react";
 
-type ReplacementFn = Array<(domNode: DOMNode) => JSX.Element | undefined>;
+type ReplacementFn = NonNullable<HTMLReactParserOptions["replace"]>;
+type TransformFn = NonNullable<HTMLReactParserOptions["transform"]>;
 
 interface Options {
   html: string;
-  additionalReplacements?: ReplacementFn;
+  additionalReplacements?: ReplacementFn[];
 }
 
 const replaceValidTags = (domNode: DOMNode) => {
@@ -57,23 +63,70 @@ const replaceSrcScripts = (domNode: DOMNode): JSX.Element | undefined => {
   }
 };
 
+// Add copy button to codehilite blocks
+const addCopyButtonToCodehilite: TransformFn = (
+  reactNode: ReactNode,
+  domNode: DOMNode,
+): JSX.Element | undefined => {
+  if (
+    domNode instanceof Element &&
+    domNode.name === "div" &&
+    domNode.attribs?.class?.includes("codehilite")
+  ) {
+    return <CopyableCode>{reactNode}</CopyableCode>;
+  }
+};
+
+const CopyableCode = ({ children }: { children: ReactNode }) => {
+  const id = useId();
+  return (
+    <div className="relative group" id={id}>
+      {children}
+
+      <CopyClipboardIcon
+        tooltip={false}
+        className="absolute top-2 right-2 p-1 opacity-0 group-hover:opacity-100 transition-opacity"
+        value={() => {
+          const codeElement = document.getElementById(id)?.firstChild;
+          if (codeElement) {
+            return codeElement.textContent || "";
+          }
+          return "";
+        }}
+      />
+    </div>
+  );
+};
+
 export const renderHTML = ({ html, additionalReplacements = [] }: Options) => {
-  const renderFunctions: ReplacementFn = [
+  const renderFunctions: ReplacementFn[] = [
     replaceValidTags,
     replaceValidIframes,
     replaceSrcScripts,
     ...additionalReplacements,
   ];
 
+  const transformFunctions: TransformFn[] = [addCopyButtonToCodehilite];
+
   return parse(html, {
-    replace: (domNode: DOMNode) => {
+    replace: (domNode: DOMNode, index: number) => {
       for (const renderFunction of renderFunctions) {
-        const replacement = renderFunction(domNode);
+        const replacement = renderFunction(domNode, index);
         if (replacement) {
           return replacement;
         }
       }
       return domNode;
     },
+    transform: (reactNode: ReactNode, domNode: DOMNode, index: number) => {
+      for (const transformFunction of transformFunctions) {
+        const transformed = transformFunction(reactNode, domNode, index);
+        if (transformed) {
+          return transformed;
+        }
+      }
+      // eslint-disable-next-line react/jsx-no-useless-fragment
+      return reactNode as JSX.Element;
+    },
   });
 };

frontend/src/plugins/core/__test__/RenderHTML.test.ts

@@ -52,6 +52,27 @@ describe("RenderHTML", () => {
     `);
   });
 
+  test("codehilite with copy button", () => {
+    const html =
+      '<div class="codehilite"><pre><code>console.log("Hello");</code></pre></div>';
+    const result = renderHTML({ html });
+
+    // Check that the result is wrapped in a CopyableCode component
+    expect(result).toMatchInlineSnapshot(`
+      <CopyableCode>
+        <div
+          className="codehilite"
+        >
+          <pre>
+            <code>
+              console.log("Hello");
+            </code>
+          </pre>
+        </div>
+      </CopyableCode>
+    `);
+  });
+
   test("custom tags - valid", () => {
     let html = "<foobar></foobar>";
     expect(renderHTML({ html })).toMatchInlineSnapshot("<foobar />");

marimo/_smoke_tests/markdown/codeblocks.py

@@ -0,0 +1,63 @@
+import marimo
+
+__generated_with = "0.11.17"
+app = marimo.App(width="medium")
+
+
+@app.cell
+def _(mo):
+    mo.md(r"""# hello""")
+    return
+
+
+@app.cell
+def _(mo):
+    mo.md(r"""This is `inline code`""")
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ```python
+        def fibonacci(n):
+            if n <= 0:
+                return []
+            elif n == 1:
+                return [0]
+            elif n == 2:
+                return [0, 1]
+
+            fib_sequence = [0, 1]
+            for i in range(2, n):
+                fib_sequence.append(fib_sequence[i-1] + fib_sequence[i-2])
+
+            return fib_sequence
+        ```
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ```
+        # Example usage:
+        # print(fibonacci(10))
+        ```
+        """
+    )
+    return
+
+
+@app.cell
+def _():
+    import marimo as mo
+    return (mo,)
+
+
+if __name__ == "__main__":
+    app.run()