UI: Add support for expand/collapse keyboard shortcuts

lib/ui/src/components/sidebar/useExpanded.ts

@@ -1,4 +1,5 @@
-import { StoriesHash } from '@storybook/api';
+import { StoriesHash, useStorybookApi } from '@storybook/api';
+import { STORIES_COLLAPSE_ALL, STORIES_EXPAND_ALL } from '@storybook/core-events';
 import { document } from 'global';
 import throttle from 'lodash/throttle';
 import React, { Dispatch, MutableRefObject, useCallback, useEffect, useReducer } from 'react';
@@ -46,6 +47,8 @@ const initializeExpanded = ({
   );
 };
 
+const noop = () => {};
+
 export const useExpanded = ({
   containerRef,
   isBrowsing,
@@ -57,6 +60,8 @@ export const useExpanded = ({
   selectedStoryId,
   onSelectStoryId,
 }: ExpandedProps): [Record<string, boolean>, Dispatch<ExpandAction>] => {
+  const api = useStorybookApi();
+
   // Track the set of currently expanded nodes within this tree.
   // Root nodes are expanded by default (and cannot be collapsed).
   const [expanded, setExpanded] = useReducer<
@@ -106,6 +111,27 @@ export const useExpanded = ({
     setExpanded({ ids: getAncestorIds(data, selectedStoryId), value: true });
   }, [data, selectedStoryId]);
 
+  const collapseAll = useCallback(() => {
+    const ids = Object.keys(data).filter(id => !rootIds.includes(id));
+    setExpanded({ ids, value: false });
+  }, [data, rootIds]);
+
+  const expandAll = useCallback(() => {
+    setExpanded({ ids: Object.keys(data), value: true });
+  }, [data]);
+
+  useEffect(() => {
+    if (!api) return noop;
+
+    api.on(STORIES_COLLAPSE_ALL, collapseAll);
+    api.on(STORIES_EXPAND_ALL, expandAll);
+
+    return () => {
+      api.off(STORIES_COLLAPSE_ALL, collapseAll);
+      api.off(STORIES_EXPAND_ALL, expandAll);
+    };
+  }, [api, collapseAll, expandAll]);
+
   // Expand, collapse or select nodes in the tree using keyboard shortcuts.
   useEffect(() => {
     const menuElement = document.getElementById('storybook-explorer-menu');