feat: add ability to Expand/Collapse All

src/__stories__/Default.tsx

@@ -7,7 +7,9 @@ import { JsonSchemaProps, JsonSchemaViewer } from '../components/JsonSchemaViewe
 
 const defaultSchema = require('../__fixtures__/default-schema.json');
 const stressSchema = require('../__fixtures__/stress-schema.json');
+const githubIssueSchema = require('../__fixtures__/real-world/github-issue.json');
 const arrayOfComplexObjects = require('../__fixtures__/arrays/of-complex-objects.json');
+const allOfComplexSchema = require('../__fixtures__/combiners/allOfs/complex.json');
 
 export default {
   component: JsonSchemaViewer,
@@ -22,12 +24,16 @@ export const Default = Template.bind({});
 
 export const CustomRowAddon = Template.bind({});
 CustomRowAddon.args = {
-  renderRowAddon: () => (
-    <Flex h="full" alignItems="center">
-      <Button pl={1} mr={1} size="sm" appearance="minimal" icon="bullseye" />
-      <input type="checkbox" />
-    </Flex>
-  ),
+  renderRowAddon: ({ nestingLevel }) => {
+    return nestingLevel == 1 ? (
+      <Flex h="full" alignItems="center">
+        <Button pl={1} mr={1} size="sm" appearance="minimal" icon="bullseye" />
+        <Button pl={1} mr={1} size="sm" appearance="minimal">
+          Expand All
+        </Button>
+      </Flex>
+    ) : null;
+  },
 };
 
 export const Expansions = Template.bind({});
@@ -88,3 +94,17 @@ export const DarkMode: Story<JsonSchemaProps> = ({ schema = defaultSchema as JSO
     </Box>
   </InvertTheme>
 );
+
+export const ExpandCollapseAll = Template.bind({});
+ExpandCollapseAll.args = {
+  showExpandAll: true,
+  schema: githubIssueSchema as JSONSchema4,
+};
+ExpandCollapseAll.storyName = 'Expand/Collapse All';
+
+export const CircularExpandCollapseAll = Template.bind({});
+CircularExpandCollapseAll.args = {
+  showExpandAll: true,
+  maxRefDepth: 10,
+  schema: allOfComplexSchema as JSONSchema4,
+};

src/components/JsonSchemaViewer.tsx

@@ -1,22 +1,22 @@
 import {
   isRegularNode,
   RootNode,
+  SchemaNodeKind,
   SchemaTree as JsonSchemaTree,
   SchemaTreeRefDereferenceFn,
 } from '@stoplight/json-schema-tree';
-import { Box, Provider as MosaicProvider } from '@stoplight/mosaic';
+import { Box, Button, HStack, Provider as MosaicProvider } from '@stoplight/mosaic';
 import { ErrorBoundaryForwardedProps, FallbackProps, withErrorBoundary } from '@stoplight/react-error-boundary';
 import cn from 'classnames';
-import { Provider } from 'jotai';
-import { useUpdateAtom } from 'jotai/utils';
+import { Provider, useAtom, useSetAtom } from 'jotai';
 import * as React from 'react';
 
 import { JSVOptions, JSVOptionsContextProvider } from '../contexts';
-import { shouldNodeBeIncluded } from '../tree/utils';
+import { isNonEmptyParentNode, shouldNodeBeIncluded } from '../tree/utils';
 import { JSONSchema } from '../types';
 import { PathCrumbs } from './PathCrumbs';
 import { TopLevelSchemaRow } from './SchemaRow';
-import { hoveredNodeAtom } from './SchemaRow/state';
+import { ExpansionMode, expansionModeAtom, hoveredNodeAtom } from './SchemaRow/state';
 
 export type JsonSchemaProps = Partial<JSVOptions> & {
   schema: JSONSchema;
@@ -29,11 +29,14 @@ export type JsonSchemaProps = Partial<JSVOptions> & {
   maxHeight?: number;
   parentCrumbs?: string[];
   skipTopLevelDescription?: boolean;
+  showExpandAll?: boolean;
 };
 
 const JsonSchemaViewerComponent = ({
   viewMode = 'standalone',
   defaultExpandedDepth = 1,
+  maxRefDepth,
+  showExpandAll = true,
   onGoToRef,
   renderRowAddon,
   renderExtensionAddon,
@@ -47,7 +50,9 @@ const JsonSchemaViewerComponent = ({
   const options = React.useMemo(
     () => ({
       defaultExpandedDepth,
+      maxRefDepth,
       viewMode,
+      showExpandAll,
       onGoToRef,
       renderRowAddon,
       renderExtensionAddon,
@@ -58,7 +63,9 @@ const JsonSchemaViewerComponent = ({
     }),
     [
       defaultExpandedDepth,
+      maxRefDepth,
       viewMode,
+      showExpandAll,
       onGoToRef,
       renderRowAddon,
       renderExtensionAddon,
@@ -73,7 +80,12 @@ const JsonSchemaViewerComponent = ({
     <MosaicProvider>
       <JSVOptionsContextProvider value={options}>
         <Provider>
-          <JsonSchemaViewerInner viewMode={viewMode} skipTopLevelDescription={skipTopLevelDescription} {...rest} />
+          <JsonSchemaViewerInner
+            viewMode={viewMode}
+            showExpandAll={showExpandAll}
+            skipTopLevelDescription={skipTopLevelDescription}
+            {...rest}
+          />
         </Provider>
       </JSVOptionsContextProvider>
     </MosaicProvider>
@@ -83,6 +95,7 @@ const JsonSchemaViewerComponent = ({
 const JsonSchemaViewerInner = ({
   schema,
   viewMode,
+  showExpandAll,
   className,
   resolveRef,
   maxRefDepth,
@@ -95,6 +108,7 @@ const JsonSchemaViewerInner = ({
   JsonSchemaProps,
   | 'schema'
   | 'viewMode'
+  | 'showExpandAll'
   | 'className'
   | 'resolveRef'
   | 'maxRefDepth'
@@ -104,11 +118,18 @@ const JsonSchemaViewerInner = ({
   | 'parentCrumbs'
   | 'skipTopLevelDescription'
 >) => {
-  const setHoveredNode = useUpdateAtom(hoveredNodeAtom);
+  const setHoveredNode = useSetAtom(hoveredNodeAtom);
+  const [expansionMode, setExpansionMode] = useAtom(expansionModeAtom);
+
   const onMouseLeave = React.useCallback(() => {
     setHoveredNode(null);
   }, [setHoveredNode]);
 
+  const onCollapseExpandAll = React.useCallback(() => {
+    const newExpansionMode: ExpansionMode = expansionMode === 'expand_all' ? 'collapse_all' : 'expand_all';
+    setExpansionMode(newExpansionMode);
+  }, [expansionMode, setExpansionMode]);
+
   const { jsonSchemaTreeRoot, nodeCount } = React.useMemo(() => {
     const jsonSchemaTree = new JsonSchemaTree(schema, {
       mergeAllOf: true,
@@ -123,6 +144,7 @@ const JsonSchemaViewerInner = ({
         nodeCount++;
         return true;
       }
+
       return false;
     });
     jsonSchemaTree.populate();
@@ -144,6 +166,21 @@ const JsonSchemaViewerInner = ({
     () => jsonSchemaTreeRoot.children.every(node => !isRegularNode(node) || node.unknown),
     [jsonSchemaTreeRoot],
   );
+
+  // Naive check if there are collapsible rows
+  const hasCollapsibleRows = React.useMemo(() => {
+    if (jsonSchemaTreeRoot.children.length === 0) {
+      return false;
+    }
+
+    if (isNonEmptyParentNode(jsonSchemaTreeRoot.children[0])) {
+      return jsonSchemaTreeRoot.children[0].children.some(childNode => {
+        return isRegularNode(childNode) && childNode.primaryType === SchemaNodeKind.Object;
+      });
+    }
+    return false;
+  }, [jsonSchemaTreeRoot]);
+
   if (isEmpty) {
     return (
       <Box className={cn(className, 'JsonSchemaViewer')} fontSize="sm" data-test="empty-text">
@@ -160,7 +197,15 @@ const JsonSchemaViewerInner = ({
       onMouseLeave={onMouseLeave}
       style={{ maxHeight }}
     >
-      <PathCrumbs parentCrumbs={parentCrumbs} />
+      {hasCollapsibleRows && showExpandAll ? (
+        <HStack w="full" fontFamily="mono" fontSize="sm" lineHeight="none" bg="canvas-pure" px="px" color="light">
+          <Box flex={1}>&nbsp;</Box>
+          <Button pl={1} mr={1} size="sm" appearance="minimal" onClick={onCollapseExpandAll}>
+            {expansionMode === 'expand_all' ? 'Collapse All' : 'Expand All'}
+          </Button>
+        </HStack>
+      ) : null}
+      <PathCrumbs parentCrumbs={parentCrumbs} showExpandAll={showExpandAll} />
       <TopLevelSchemaRow schemaNode={jsonSchemaTreeRoot.children[0]} skipDescription={skipTopLevelDescription} />
     </Box>
   );

src/components/PathCrumbs/index.tsx

@@ -1,15 +1,28 @@
-import { Box, HStack } from '@stoplight/mosaic';
+import { Box, Button, HStack } from '@stoplight/mosaic';
 import { useAtom } from 'jotai';
 import * as React from 'react';
 
 import { useJSVOptionsContext } from '../../contexts';
+import { ExpansionMode, expansionModeAtom } from '../SchemaRow/state';
 import { pathCrumbsAtom, showPathCrumbsAtom } from './state';
 
-export const PathCrumbs = ({ parentCrumbs = [] }: { parentCrumbs?: string[] }) => {
+export const PathCrumbs = ({
+  parentCrumbs = [],
+  showExpandAll = false,
+}: {
+  parentCrumbs?: string[];
+  showExpandAll?: boolean;
+}) => {
   const [showPathCrumbs] = useAtom(showPathCrumbsAtom);
   const [pathCrumbs] = useAtom(pathCrumbsAtom);
+  const [expansionMode, setExpansionMode] = useAtom(expansionModeAtom);
   const { disableCrumbs } = useJSVOptionsContext();
 
+  const onCollapseExpandAll = React.useCallback(() => {
+    const newExpansionMode: ExpansionMode = expansionMode === 'expand_all' ? 'collapse_all' : 'expand_all';
+    setExpansionMode(newExpansionMode);
+  }, [expansionMode, setExpansionMode]);
+
   if (disableCrumbs) {
     return null;
   }
@@ -39,8 +52,6 @@ export const PathCrumbs = ({ parentCrumbs = [] }: { parentCrumbs?: string[] }) =
 
   return (
     <HStack
-      spacing={1}
-      divider={<Box>/</Box>}
       h="md"
       // so that the crumbs take up no space in the dom, and thus do not push content down when they appear
       mt={-8}
@@ -56,8 +67,31 @@ export const PathCrumbs = ({ parentCrumbs = [] }: { parentCrumbs?: string[] }) =
       color="light"
       alignItems="center"
     >
-      {parentCrumbElems}
-      {pathCrumbElems.length && <HStack divider={<Box fontWeight="bold">.</Box>}>{pathCrumbElems}</HStack>}
+      <HStack
+        spacing={1}
+        divider={<Box>/</Box>}
+        w="full"
+        h="md"
+        borderB
+        fontFamily="mono"
+        fontSize="sm"
+        lineHeight="none"
+        bg="canvas-pure"
+        px="px"
+        color="light"
+        alignItems="center"
+        justifyContent="between"
+      >
+        {parentCrumbElems}
+        {pathCrumbElems.length && <HStack divider={<Box fontWeight="bold">.</Box>}>{pathCrumbElems}</HStack>}
+      </HStack>
+      {showExpandAll ? (
+        <Box flex={1}>
+          <Button pl={1} mr={1} size="sm" appearance="minimal" onClick={onCollapseExpandAll}>
+            {expansionMode === 'expand_all' ? 'Collapse All' : 'Expand All'}
+          </Button>
+        </Box>
+      ) : null}
     </HStack>
   );
 };

src/components/SchemaRow/SchemaRow.tsx

@@ -1,8 +1,7 @@
 import { isMirroredNode, isReferenceNode, isRegularNode, SchemaNode } from '@stoplight/json-schema-tree';
 import { Box, Flex, NodeAnnotation, Select, SpaceVals, VStack } from '@stoplight/mosaic';
 import type { ChangeType } from '@stoplight/types';
-import { Atom } from 'jotai';
-import { useAtomValue, useUpdateAtom } from 'jotai/utils';
+import { Atom, useAtomValue, useSetAtom } from 'jotai';
 import last from 'lodash/last.js';
 import * as React from 'react';
 
@@ -15,7 +14,7 @@ import { Caret, Description, getValidationsFromSchema, Types, Validations } from
 import { ChildStack } from '../shared/ChildStack';
 import { Error } from '../shared/Error';
 import { Properties, useHasProperties } from '../shared/Properties';
-import { hoveredNodeAtom, isNodeHoveredAtom } from './state';
+import { expansionModeAtom, hoveredNodeAtom, isNodeHoveredAtom } from './state';
 import { useChoices } from './useChoices';
 
 export interface SchemaRowProps {
@@ -30,6 +29,7 @@ export const SchemaRow: React.FunctionComponent<SchemaRowProps> = React.memo(
   ({ schemaNode, nestingLevel, pl, parentNodeId, parentChangeType }) => {
     const {
       defaultExpandedDepth,
+      maxRefDepth,
       renderRowAddon,
       renderExtensionAddon,
       onGoToRef,
@@ -39,7 +39,9 @@ export const SchemaRow: React.FunctionComponent<SchemaRowProps> = React.memo(
       viewMode,
     } = useJSVOptionsContext();
 
-    const setHoveredNode = useUpdateAtom(hoveredNodeAtom);
+    const setHoveredNode = useSetAtom(hoveredNodeAtom);
+
+    const expansionMode = useAtomValue(expansionModeAtom);
 
     const nodeId = getNodeId(schemaNode, parentNodeId);
 
@@ -49,7 +51,9 @@ export const SchemaRow: React.FunctionComponent<SchemaRowProps> = React.memo(
     const hasChanged = nodeHasChanged?.({ nodeId: originalNodeId, mode });
 
     const [isExpanded, setExpanded] = React.useState<boolean>(
-      !isMirroredNode(schemaNode) && nestingLevel <= defaultExpandedDepth,
+      expansionMode === 'expand_all'
+        ? !isMirroredNode(schemaNode)
+        : !isMirroredNode(schemaNode) && nestingLevel <= defaultExpandedDepth,
     );
 
     const { selectedChoice, setSelectedChoice, choices } = useChoices(schemaNode);
@@ -67,6 +71,17 @@ export const SchemaRow: React.FunctionComponent<SchemaRowProps> = React.memo(
     const validations = isRegularNode(schemaNode) ? schemaNode.validations : {};
     const hasProperties = useHasProperties({ required, deprecated, validations });
 
+    React.useEffect(() => {
+      if (expansionMode === 'expand_all' && !isExpanded) {
+        const canBeExpanded = maxRefDepth && maxRefDepth > 0 ? nestingLevel < maxRefDepth : true;
+        if (canBeExpanded) {
+          setExpanded(true);
+        }
+      } else if (expansionMode === 'collapse_all' && isExpanded) {
+        setExpanded(false);
+      }
+    }, [isExpanded, expansionMode, nestingLevel, maxRefDepth]);
+
     const [totalVendorExtensions, vendorExtensions] = React.useMemo(
       () => extractVendorExtensions(schemaNode.fragment),
       [schemaNode.fragment],

src/components/SchemaRow/TopLevelSchemaRow.tsx

@@ -1,7 +1,7 @@
 import { isPlainObject } from '@stoplight/json';
 import { isRegularNode, RegularNode } from '@stoplight/json-schema-tree';
 import { Box, Flex, HStack, Icon, Menu, Pressable } from '@stoplight/mosaic';
-import { useUpdateAtom } from 'jotai/utils';
+import { useSetAtom } from 'jotai';
 import { isEmpty } from 'lodash';
 import * as React from 'react';
 
@@ -22,7 +22,6 @@ export const TopLevelSchemaRow = ({
   skipDescription,
 }: Pick<SchemaRowProps, 'schemaNode'> & { skipDescription?: boolean }) => {
   const { renderExtensionAddon } = useJSVOptionsContext();
-
   const { selectedChoice, setSelectedChoice, choices } = useChoices(schemaNode);
   const childNodes = React.useMemo(() => visibleChildren(selectedChoice.type), [selectedChoice.type]);
   const nestingLevel = 0;
@@ -150,7 +149,7 @@ function ScrollCheck() {
   const elementRef = React.useRef<HTMLDivElement>(null);
 
   const isOnScreen = useIsOnScreen(elementRef);
-  const setShowPathCrumbs = useUpdateAtom(showPathCrumbsAtom);
+  const setShowPathCrumbs = useSetAtom(showPathCrumbsAtom);
   React.useEffect(() => {
     setShowPathCrumbs(!isOnScreen);
   }, [isOnScreen, setShowPathCrumbs]);

src/components/SchemaRow/state.ts

@@ -2,6 +2,9 @@ import { SchemaNode } from '@stoplight/json-schema-tree';
 import { atom } from 'jotai';
 import { atomFamily } from 'jotai/utils';
 
+export type ExpansionMode = 'expand_all' | 'collapse_all' | 'off';
+
+export const expansionModeAtom = atom<ExpansionMode>('off');
 export const hoveredNodeAtom = atom<SchemaNode | null>(null);
 export const isNodeHoveredAtom = atomFamily((node: SchemaNode) => atom(get => node === get(hoveredNodeAtom)));
 export const isChildNodeHoveredAtom = atomFamily((parent: SchemaNode) =>