Refactor modular pipeline expansions to instantly update UI

RELEASE.md

@@ -14,6 +14,7 @@ Please follow the established format:
 ## Bug fixes and other changes
 
 - Fix kedro viz `--load-file` to run from any directory without requiring a Kedro project. (#2206)
+- Improved modular pipeline expand/collapse logic for better state synchronisation. (#2225)
 
 # Release 10.1.0
 

src/actions/index.js

@@ -27,7 +27,7 @@ export function toggleLayers(visible) {
 export const TOGGLE_EXPAND_ALL_PIPELINES = 'TOGGLE_EXPAND_ALL_PIPELINES';
 
 /**
- * Toggle whether to expand all modular pipelines or collapse
+ * Toggle whether to expand all modular pipelines or collapse with boolean.
  * @param {Boolean} shouldExpandAllPipelines
  */
 export function toggleExpandAllPipelines(shouldExpandAllPipelines) {

src/actions/modular-pipelines.js

@@ -55,3 +55,17 @@ export function toggleModularPipelinesExpanded(expandedIDs) {
     expandedIDs,
   };
 }
+
+export const TOGGLE_MODULAR_PIPELINES_VISIBILITY_STATE =
+  'TOGGLE_MODULAR_PIPELINES_VISIBILITY_STATE';
+
+/**
+ * Toggles the visibility state for all modular pipelines (expand or collapse all).
+ * @param {Boolean} expandAllPipelines - Whether to expand (true) or collapse (false) all pipelines.
+ */
+export function toggleModularPipelinesVisibilityState(expandAllPipelines) {
+  return {
+    type: TOGGLE_MODULAR_PIPELINES_VISIBILITY_STATE,
+    expandAllPipelines,
+  };
+}

src/components/flowchart-primary-toolbar/flowchart-primary-toolbar.js

@@ -7,7 +7,7 @@ import {
   toggleTextLabels,
   toggleExpandAllPipelines,
 } from '../../actions';
-import { loadInitialPipelineData } from '../../actions/pipelines';
+import { toggleModularPipelinesVisibilityState } from '../../actions/modular-pipelines';
 import IconButton from '../ui/icon-button';
 import LabelIcon from '../icons/label';
 import ExportIcon from '../icons/export';
@@ -126,7 +126,7 @@ export const mapDispatchToProps = (dispatch) => ({
   },
   onToggleExpandAllPipelines: (isExpanded) => {
     dispatch(toggleExpandAllPipelines(isExpanded));
-    dispatch(loadInitialPipelineData());
+    dispatch(toggleModularPipelinesVisibilityState(isExpanded));
   },
 });
 

src/components/flowchart-primary-toolbar/flowchart-primary-toolbar.test.js

@@ -59,6 +59,8 @@ describe('PrimaryToolbar', () => {
         visible: mockState.spaceflights.visible,
         display: mockState.spaceflights.display,
         [callback]: mockFn,
+        modularPipelineIDs: ['pipeline1', '__root__'],
+        onToggleExpandPipelines: jest.fn(),
       };
       const wrapper = setup.mount(<FlowchartPrimaryToolbar {...props} />);
       expect(mockFn.mock.calls.length).toBe(0);

src/reducers/modular-pipelines.js

@@ -3,6 +3,7 @@ import {
   TOGGLE_MODULAR_PIPELINES_EXPANDED,
   TOGGLE_SINGLE_MODULAR_PIPELINE_EXPANDED,
   TOGGLE_MODULAR_PIPELINE_DISABLED,
+  TOGGLE_MODULAR_PIPELINES_VISIBILITY_STATE,
 } from '../actions/modular-pipelines';
 
 function modularPipelineReducer(modularPipelineState = {}, action) {
@@ -110,6 +111,37 @@ function modularPipelineReducer(modularPipelineState = {}, action) {
         visible: newVisibleState,
       });
     }
+    case TOGGLE_MODULAR_PIPELINES_VISIBILITY_STATE: {
+      let newVisibleState = {};
+
+      // Determine which IDs should be expanded based on the action
+      const expandedIDs = action.expandAllPipelines
+        ? modularPipelineState.ids
+        : [];
+
+      if (action.expandAllPipelines) {
+        // If expanding all pipelines, set visibility for each pipeline and its children
+        modularPipelineState.ids.forEach((modularPipelineID) => {
+          newVisibleState[modularPipelineID] = false;
+          modularPipelineState.tree[modularPipelineID].children.forEach(
+            (child) => (newVisibleState[child.id] = true)
+          );
+        });
+      } else {
+        // If not expanding all, only set visibility for the children of the root pipeline
+        if (modularPipelineState.tree['__root__']) {
+          for (const child of modularPipelineState.tree['__root__'].children ||
+            []) {
+            newVisibleState[child.id] = true;
+          }
+        }
+      }
+
+      return updateState({
+        expanded: expandedIDs,
+        visible: newVisibleState,
+      });
+    }
 
     default:
       return modularPipelineState;