feat: add expand and collapse all columns to notebook actions menu

frontend/src/components/editor/actions/useNotebookActions.tsx

@@ -38,6 +38,8 @@ import {
   MessagesSquareIcon,
   YoutubeIcon,
   DiamondPlusIcon,
+  ChevronDownCircleIcon,
+  ChevronRightCircleIcon,
 } from "lucide-react";
 import { commandPaletteAtom } from "../controls/command-palette";
 import {
@@ -81,6 +83,7 @@ import { useRunAllCells } from "../cell/useRunCells";
 import { settingDialogAtom } from "@/components/app-config/state";
 import { AddDatabaseDialogContent } from "../database/add-database-form";
 import { useHideAllMarkdownCode } from "./useHideAllMarkdownCode";
+import { useSectionCollapse } from "./useSectionCollapse";
 import { Constants } from "@/core/constants";
 import { getFeatureFlag } from "@/core/config/feature-flag";
 
@@ -97,6 +100,7 @@ export function useNotebookActions() {
   const [viewState] = useAtom(viewStateAtom);
   const kioskMode = useAtomValue(kioskModeAtom);
   const hideAllMarkdownCode = useHideAllMarkdownCode();
+  const { collapseAllSections, expandAllSections } = useSectionCollapse();
 
   const {
     updateCellConfig,
@@ -375,6 +379,16 @@ export function useNotebookActions() {
       label: "Hide all markdown code",
       handle: hideAllMarkdownCode,
     },
+    {
+      icon: <ChevronRightCircleIcon size={14} strokeWidth={1.5} />,
+      label: "Collapse all sections",
+      handle: collapseAllSections,
+    },
+    {
+      icon: <ChevronDownCircleIcon size={14} strokeWidth={1.5} />,
+      label: "Expand all sections",
+      handle: expandAllSections,
+    },
     {
       icon: <DiamondPlusIcon size={14} strokeWidth={1.5} />,
       label: "Add setup cell",

frontend/src/components/editor/actions/useSectionCollapse.ts

@@ -0,0 +1,38 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { useCellActions } from "@/core/cells/cells";
+import { getNotebook } from "@/core/cells/cells";
+import { canCollapseOutline } from "@/core/dom/outline";
+import useEvent from "react-use-event-hook";
+
+/**
+ * Hooks to collapse and expand all sections in the notebook.
+ */
+
+export const useSectionCollapse = () => {
+  const { collapseCell, expandCell } = useCellActions();
+
+  const processAllSections = async (action: "collapse" | "expand") => {
+    const notebook = getNotebook();
+    const cellIds = notebook.cellIds.inOrderIds;
+
+    // Find all markdown cells that aren't already hidden
+    for (const cellId of cellIds) {
+      const outline = notebook.cellRuntime[cellId].outline;
+      // Check if the cell is a markdown cell with a TOC outline
+      if (!outline) {
+        continue;
+      }
+      // Check if the cell is a collapsible header
+      if (!canCollapseOutline(outline)) {
+        continue;
+      }
+      // Collapse or expand the cell based on the action
+      action === "collapse" ? collapseCell({ cellId }) : expandCell({ cellId });
+    }
+  };
+
+  return {
+    collapseAllSections: useEvent(() => processAllSections("collapse")),
+    expandAllSections: useEvent(() => processAllSections("expand")),
+  };
+};