Treeview fixes (very un-optimised)

ndelangen · ndelangen · commit bf45f06cd8b9 · 2018-11-19T13:22:52.000+01:00
diff --git a/lib/components/src/explorer/explorer.js b/lib/components/src/explorer/explorer.js
@@ -1,23 +1,45 @@
 import React, { Component } from 'react';
 import PropTypes from 'prop-types';
 
+import memoize from 'memoizerific';
+
 import { TreeState } from '../treeview/treeview';
 import Placeholder from '../placeholder/placeholder';
+import { getParents } from '../treeview/utils';
+
+const toSelected = memoize(50)((id, stories) =>
+  Object.keys(stories).reduce((acc, k) => Object.assign(acc, { [k]: k === id }), {})
+);
+const toExpanded = memoize(50)((selectedIds, stories) => {
+  const parents = selectedIds.reduce(
+    (acc, id) => acc.concat(getParents({ path: id, dataset: stories }).map(i => i.path)),
+    []
+  );
+  return Object.keys(stories).reduce(
+    (acc, k) => Object.assign(acc, { [k]: parents.includes(k) }),
+    {}
+  );
+});
 
 // This component gets a ref so it needs to be a class
 // eslint-disable-next-line react/prefer-stateless-function
 class Explorer extends Component {
   render() {
-    const { stories, ...rest } = this.props;
+    const { stories, componentId, ...rest } = this.props;
     const list = Object.entries(stories);
+    const selected = toSelected(componentId, stories);
+    const expanded = toExpanded(
+      Object.entries(selected).reduce((acc, [k, v]) => (v ? acc.concat(k) : acc), []),
+      stories
+    );
 
     return list.length ? (
       <TreeState
         dataset={stories}
         prefix="explorer"
         {...rest}
-        selected={{}}
-        expanded={{}}
+        selected={selected}
+        expanded={expanded}
         filter=""
       />
     ) : (
@@ -28,6 +50,7 @@ class Explorer extends Component {
 
 Explorer.propTypes = {
   stories: PropTypes.shape({}).isRequired,
+  componentId: PropTypes.string,
 };
 
 export { Explorer };
diff --git a/lib/components/src/treeview/treeview.js b/lib/components/src/treeview/treeview.js
@@ -273,7 +273,16 @@ class TreeState extends PureComponent {
 
   render() {
     const { Head, Leaf, Branch, Filter, RootTitle, events } = this;
-    const { dataset, selected, expanded } = this.state;
+    const { dataset } = this.state;
+
+    const expanded = Object.keys(dataset).reduce(
+      (acc, k) => Object.assign(acc, { [k]: this.props.expanded[k] || this.state.expanded[k] }),
+      {}
+    );
+    const selected = Object.keys(dataset).reduce(
+      (acc, k) => Object.assign(acc, { [k]: this.props.selected[k] || this.state.selected[k] }),
+      {}
+    );
 
     const mains = getMains({ dataset });
     const { roots, others } = mains.reduce(
diff --git a/lib/ui/src/components/nav/explorer.js b/lib/ui/src/components/nav/explorer.js
@@ -22,7 +22,7 @@ const LeafLink = ({ id, href, children, className, ...rest }) => (
   </Router.Location>
 );
 
-const StoriesPanel = ({ stories }) => <Explorer allowClick stories={stories} Link={LeafLink} />;
+const StoriesPanel = props => <Explorer allowClick {...props} Link={LeafLink} />;
 StoriesPanel.propTypes = {
   stories: PropTypes.shape({}).isRequired,
 };
diff --git a/lib/ui/src/components/nav/nav.js b/lib/ui/src/components/nav/nav.js
@@ -78,15 +78,15 @@ const Brand = withCSSContext(({ title, url }, { theme: { brand: UserBrand } }) =
 ));
 Brand.displayName = 'Brand';
 
-const Nav = ({ title, url, notifications = [], stories, menu: items }) => (
+const Nav = ({ title, url, componentId, notifications = [], stories, menu: items }) => (
   <Container>
     <Inner>
       <Main>
         <Head>
           <Brand title={title} url={url} />
           <MenuToggle id="storybook-explorer-menu">{createMenu(items)}</MenuToggle>
         </Head>
-        <Explorer stories={stories} />
+        <Explorer stories={stories} componentId={componentId} />
       </Main>
       <Foot>
         {notifications.map(({ id }) => (
@@ -130,6 +130,7 @@ Nav.propTypes = {
   title: PropTypes.string.isRequired,
   url: PropTypes.string.isRequired,
   stories: PropTypes.shape({}).isRequired,
+  componentId: PropTypes.string,
   notifications: PropTypes.arrayOf(PropTypes.shape({})).isRequired,
   menu: PropTypes.arrayOf(
     PropTypes.shape({
diff --git a/lib/ui/src/containers/nav.js b/lib/ui/src/containers/nav.js
@@ -12,6 +12,8 @@ export const mapper = ({ state, api }) => {
   const {
     uiOptions: { name, url },
     notifications,
+    viewMode,
+    componentId,
     ui: { isFullscreen, showPanel, showNav, panelPosition },
   } = state;
   const shortcutKeys = get('shortcutKeys') || defaultKeyboardShortcuts;
@@ -21,6 +23,8 @@ export const mapper = ({ state, api }) => {
     url,
     notifications,
     stories: state.storiesHash,
+    componentId,
+    viewMode,
     menu: [
       {
         id: 'about',
