Expand All and Collapse All functionality.

Author: eugenezinovyev

packages/components/TreeView/TreeView.js

@@ -1,12 +1,13 @@
 import React, { useMemo } from 'react';
 import buildTableRowsData from '../render/buildTableRowsData';
+import LinkButton from '../LinkButton/LinkButton';
 import renderRow from '../render/renderRow';
 import renderTable from '../render/renderTable';
 import useCollapse from './useCollapse';
 import classes from './TreeView.module.less';
 
 const TreeView = ({ coverage, linkSelector }) => {
-  const collapseContainer = useCollapse();
+  const [collapseContainer, onCollapseAll, onExpandAll] = useCollapse();
   const tableHTML = useMemo(() => {
     const rows = buildTableRowsData(coverage);
     const bodyContent = rows.map(row => renderRow(row, { linkSelector }));
@@ -15,7 +16,15 @@ const TreeView = ({ coverage, linkSelector }) => {
   }, [coverage, linkSelector]);
 
 
-  return <div className={classes.root} dangerouslySetInnerHTML={tableHTML} ref={collapseContainer}/>;
+  return (
+    <div className={classes.root}>
+      <div className={classes.controls}>
+        <LinkButton onClick={onCollapseAll}>Collapse All</LinkButton>
+        <LinkButton onClick={onExpandAll}>Expand All</LinkButton>
+      </div>
+      <div dangerouslySetInnerHTML={tableHTML} ref={collapseContainer}/>
+    </div>
+  );
 };
 
 export default TreeView;

packages/components/TreeView/TreeView.module.less

@@ -1,4 +1,14 @@
 .root {
   display: flex;
   flex-direction: column;
-}
+}
+
+.controls {
+  display: flex;
+  flex-direction: row;
+  margin-bottom: 0.5rem;
+
+  & > * + * {
+    margin-left: 1rem;
+  }
+}

packages/components/TreeView/useCollapse.js

@@ -1,5 +1,5 @@
 import { BRANCH_NODE } from '@lcov-viewer/core';
-import { useEffect, useRef } from 'react';
+import { useCallback, useEffect, useRef } from 'react';
 import { collapseRow, expandRow, hideRow, showRow } from '../render/collapse';
 
 const findRow = (element) => {
@@ -25,7 +25,8 @@ const useCollapse = () => {
 
     Array.prototype.forEach.call(ref.current.querySelectorAll('tbody tr'), row => {
       const path = row.getAttribute('data-node');
-      state.current[path] = { path, row, collapsed: false, visible: true };
+      const type = row.getAttribute('data-node-type');
+      state.current[path] = { path, type, row, collapsed: false, visible: true };
     });
 
     const handleClick = (event) => {
@@ -53,8 +54,10 @@ const useCollapse = () => {
       affectedChildrenPaths.forEach(path => {
         if (rowState.collapsed) {
           showRow(state.current[path].row);
+          state.current[path].visible = true;
         } else {
           hideRow(state.current[path].row);
+          state.current[path].visible = false;
         }
       });
 
@@ -69,7 +72,38 @@ const useCollapse = () => {
     };
   }, []);
 
-  return ref;
+  const onCollapseAll = useCallback(() => {
+    Object.values(state.current).forEach(rowState => {
+      if (rowState.type === BRANCH_NODE) {
+        if (!rowState.collapsed) {
+          collapseRow(rowState.row);
+          rowState.collapsed = true;
+        }
+      }
+      
+      if (rowState.path !== '') {
+        hideRow(rowState.row);
+        rowState.visible = false;
+      }
+    });
+      
+  }, []);
+
+  const onExpandAll = useCallback(() => {
+    Object.values(state.current).forEach(rowState => {
+      if (rowState.collapsed) {
+        expandRow(rowState.row);
+        rowState.collapsed = false;
+      }
+
+      if (!rowState.visible) {
+        showRow(rowState.row);
+        rowState.visible = true;
+      }
+    });
+  }, []);
+
+  return [ref, onCollapseAll, onExpandAll];
 };
 
 export default useCollapse;