Typescript: Improve @storybook/ui types

lib/api/src/lib/stories.ts

@@ -37,7 +37,7 @@ export interface Story {
   name: string;
   kind: string;
   children?: StoryId[];
-  isComponent: false;
+  isComponent: boolean;
   isRoot: false;
   isLeaf: true;
   parameters?: {

lib/ui/src/components/sidebar/treeview/utils.ts

@@ -81,7 +81,7 @@ export function isStory(item: Item): item is Story {
 export const get = memoize(1000)((id: string, dataset: Dataset) => dataset[id]);
 export const getParent = memoize(1000)((id: string, dataset: Dataset) => {
   const item = get(id, dataset);
-  if (!isRoot(item)) {
+  if (item && !isRoot(item)) {
     return get(item.parent, dataset);
   }
   return undefined;

lib/ui/src/containers/sidebar.test.ts

@@ -0,0 +1,206 @@
+import { State } from '@storybook/api';
+import { collapseDocsOnlyStories, collapseAllStories } from './sidebar';
+
+type StoriesHash = State['storiesHash'];
+type Item = StoriesHash[keyof StoriesHash];
+
+const docsOnly = { parameters: { docsOnly: true } };
+const root: Item = {
+  id: 'root',
+  name: 'root',
+  depth: 0,
+  children: ['a', 'b'],
+  isRoot: true,
+  isComponent: false,
+  isLeaf: false,
+};
+const a: Item = {
+  id: 'a',
+  name: 'a',
+  depth: 1,
+  isComponent: true,
+  isRoot: false,
+  isLeaf: false,
+  parent: 'root',
+  children: ['a1'],
+};
+const a1: Item = {
+  id: 'a1',
+  name: 'a1',
+  kind: 'a',
+  depth: 2,
+  isLeaf: true,
+  isComponent: false,
+  isRoot: false,
+  parent: 'a',
+};
+const b: Item = {
+  id: 'b',
+  name: 'b',
+  depth: 1,
+  isRoot: false,
+  isComponent: true,
+  isLeaf: false,
+  parent: 'root',
+  children: ['b1', 'b2'],
+};
+const b1: Item = {
+  id: 'b1',
+  name: 'b1',
+  kind: 'b',
+  depth: 2,
+  isLeaf: true,
+  isRoot: false,
+  isComponent: false,
+  parent: 'b',
+};
+const b2: Item = {
+  id: 'b2',
+  name: 'b2',
+  kind: 'b',
+  depth: 2,
+  isLeaf: true,
+  isRoot: false,
+  isComponent: false,
+  parent: 'b',
+};
+
+const stories: StoriesHash = { root, a, a1, b, b1, b2 };
+
+// this seems to work as intended
+describe('collapse docs-only stories', () => {
+  it('ignores normal stories', () => {
+    const filtered = collapseDocsOnlyStories(stories);
+    expect(filtered).toEqual(stories);
+  });
+
+  it('filters out docs-only stories', () => {
+    const hasDocsOnly: StoriesHash = {
+      ...stories,
+      a1: { ...a1, parameters: { ...a1.parameters, ...docsOnly.parameters } },
+    };
+    const filtered = collapseDocsOnlyStories(hasDocsOnly);
+
+    expect(filtered.root.children).not.toContain(hasDocsOnly.a.id);
+    expect(filtered.root.children).toContain(hasDocsOnly.a1.id);
+    expect(filtered.root.children).toContain(hasDocsOnly.b.id);
+
+    expect(filtered.a).not.toBeDefined();
+  });
+});
+
+describe('collapse all stories', () => {
+  it('collapses normal stories', () => {
+    const collapsed = collapseAllStories(stories);
+
+    const expected: StoriesHash = {
+      a1: {
+        id: 'a1',
+        depth: 1,
+        name: 'a',
+        kind: 'a',
+        parent: 'root',
+        children: [],
+        isRoot: false,
+        isComponent: true,
+        isLeaf: true,
+      },
+      b1: {
+        id: 'b1',
+        depth: 1,
+        name: 'b',
+        kind: 'b',
+        parent: 'root',
+        children: [],
+        isRoot: false,
+        isComponent: true,
+        isLeaf: true,
+      },
+      root: {
+        id: 'root',
+        name: 'root',
+        depth: 0,
+        children: ['a1', 'b1'],
+        isRoot: true,
+        isComponent: false,
+        isLeaf: false,
+      },
+    };
+
+    expect(collapsed).toEqual(expected);
+  });
+
+  it('collapses docs-only stories', () => {
+    const hasDocsOnly: StoriesHash = {
+      ...stories,
+      a1: { ...a1, parameters: { ...a1.parameters, ...docsOnly.parameters } },
+    };
+
+    const collapsed = collapseAllStories(hasDocsOnly);
+
+    expect(collapsed.a1).toEqual({
+      id: 'a1',
+      name: 'a',
+      kind: 'a',
+      depth: 1,
+      isComponent: true,
+      isLeaf: true,
+      isRoot: false,
+      parent: 'root',
+      children: [],
+    });
+  });
+
+  it('collapses mixtures of leaf and non-leaf children', () => {
+    const mixedRoot: Item = {
+      id: 'root',
+      name: 'root',
+      depth: 0,
+      isRoot: true,
+      isComponent: false,
+      isLeaf: false,
+      children: ['a', 'b1'],
+    };
+
+    const mixed: StoriesHash = {
+      root: mixedRoot,
+      a,
+      a1,
+      b1: { ...b1, depth: 1, parent: 'root' },
+    };
+    const collapsed = collapseAllStories(mixed);
+
+    expect(collapsed).toEqual({
+      a1: {
+        id: 'a1',
+        depth: 1,
+        name: 'a',
+        kind: 'a',
+        isRoot: false,
+        isComponent: true,
+        isLeaf: true,
+        parent: 'root',
+        children: [],
+      },
+      b1: {
+        id: 'b1',
+        name: 'b1',
+        kind: 'b',
+        depth: 1,
+        isLeaf: true,
+        isComponent: false,
+        isRoot: false,
+        parent: 'root',
+      },
+      root: {
+        id: 'root',
+        name: 'root',
+        depth: 0,
+        children: ['a1', 'b1'],
+        isRoot: true,
+        isComponent: false,
+        isLeaf: false,
+      },
+    });
+  });
+});

lib/ui/src/containers/sidebar.tsx

@@ -5,6 +5,7 @@ import memoize from 'memoizerific';
 import { Badge } from '@storybook/components';
 import { Consumer, Combo, State } from '@storybook/api';
 
+import { Group, Story } from '@storybook/api/dist/lib/stories';
 import { shortcutToHumanString } from '../libs/shortcut';
 
 import ListItemIcon from '../components/sidebar/ListItemIcon';
@@ -117,11 +118,13 @@ export const collapseAllStories = (stories: StoriesHash) => {
   const componentIdToLeafId: Record<string, string> = {};
 
   // 1) remove all leaves
-  const leavesRemoved = Object.values(stories).filter(item => isStory(item));
+  const leavesRemoved = Object.values(stories).filter(
+    item => !(item.isLeaf && stories[item.parent].isComponent)
+  );
 
   // 2) make all components leaves and rewrite their ID's to the first leaf child
   const componentsFlattened = leavesRemoved.map(item => {
-    const { id, isComponent, children, ...rest } = item;
+    const { id, isComponent, isRoot, children, ...rest } = item;
 
     // this is a folder, so just leave it alone
     if (!isComponent) {
@@ -140,6 +143,8 @@ export const collapseAllStories = (stories: StoriesHash) => {
     const component = {
       ...rest,
       id: leafId,
+      kind: (stories[leafId] as Story).kind,
+      isRoot: false,
       isLeaf: true,
       isComponent: true,
       children: [] as string[],
@@ -185,6 +190,9 @@ export const collapseDocsOnlyStories = (storiesHash: StoriesHash) => {
     }
     return true;
   });
+
+  // console.dir(docsOnlyStoriesRemoved, { depth: 10 });
+
   const docsOnlyComponentsCollapsed = docsOnlyStoriesRemoved.map(item => {
     // collapse docs-only components
     const { isComponent, children, id } = item;