sankey: implement node grouping via mouse selection

src/components/modebar/buttons.js

@@ -513,6 +513,17 @@ function toggleHover(gd) {
     Registry.call('_guiRelayout', gd, 'hovermode', newHover);
 }
 
+modeBarButtons.resetSankeyGroup = {
+    name: 'resetSankeyGroup',
+    title: function(gd) { return _(gd, 'Ungroup all nodes'); },
+    icon: Icons.home,
+    click: function(gd) {
+        Registry.call('restyle', gd, {
+            'node.groups': [[]],
+        });
+    }
+};
+
 // buttons when more then one plot types are present
 
 modeBarButtons.toggleHover = {

src/components/modebar/manage.js

@@ -86,6 +86,7 @@ function getButtonGroups(gd, buttonsToRemove, buttonsToAdd, showSendToCloud) {
     var hasTernary = fullLayout._has('ternary');
     var hasMapbox = fullLayout._has('mapbox');
     var hasPolar = fullLayout._has('polar');
+    var hasSankey = fullLayout._has('sankey');
     var allAxesFixed = areAllAxesFixed(fullLayout);
 
     var groups = [];
@@ -139,6 +140,9 @@ function getButtonGroups(gd, buttonsToRemove, buttonsToAdd, showSendToCloud) {
     else if(hasPie) {
         hoverGroup = ['hoverClosestPie'];
     }
+    else if(hasSankey) {
+        hoverGroup = ['resetSankeyGroup'];
+    }
     else { // hasPolar, hasTernary
         // always show at least one hover icon.
         hoverGroup = ['toggleHover'];

src/plots/cartesian/select.js

@@ -48,6 +48,8 @@ function prepSelect(e, startX, startY, dragOptions, mode) {
     var allAxes = dragOptions.xaxes.concat(dragOptions.yaxes);
     var subtract = e.altKey;
 
+    var doneFnCompleted = dragOptions.doneFnCompleted;
+
     var filterPoly, selectionTester, mergedPolygons, currentPolygon;
     var i, searchInfo, eventData;
 
@@ -284,6 +286,8 @@ function prepSelect(e, startX, startY, dragOptions, mode) {
                 dragOptions.mergedPolygons.length = 0;
                 [].push.apply(dragOptions.mergedPolygons, mergedPolygons);
             }
+
+            doneFnCompleted(selection);
         });
     };
 }
@@ -519,6 +523,11 @@ function determineSearchTraces(gd, xAxes, yAxes, subplot) {
             var info = createSearchInfo(trace._module, cd, xAxes[0], yAxes[0]);
             info.scene = gd._fullLayout._splomScenes[trace.uid];
             searchTraces.push(info);
+        } else if(
+          trace.type === 'sankey'
+        ) {
+            var sankeyInfo = createSearchInfo(trace._module, cd, xAxes[0], yAxes[0]);
+            searchTraces.push(sankeyInfo);
         } else {
             if(xAxisIds.indexOf(trace.xaxis) === -1) continue;
             if(yAxisIds.indexOf(trace.yaxis) === -1) continue;

src/traces/sankey/base_plot.js

@@ -13,6 +13,12 @@ var getModuleCalcData = require('../../plots/get_data').getModuleCalcData;
 var plot = require('./plot');
 var fxAttrs = require('../../components/fx/layout_attributes');
 
+var setCursor = require('../../lib/setcursor');
+var dragElement = require('../../components/dragelement');
+var prepSelect = require('../../plots/cartesian/select').prepSelect;
+var Lib = require('../../lib');
+var Registry = require('../../registry');
+
 var SANKEY = 'sankey';
 
 exports.name = SANKEY;
@@ -24,6 +30,7 @@ exports.baseLayoutAttrOverrides = overrideAll({
 exports.plot = function(gd) {
     var calcData = getModuleCalcData(gd.calcdata, SANKEY)[0];
     plot(gd, calcData);
+    exports.updateFx(gd);
 };
 
 exports.clean = function(newFullData, newFullLayout, oldFullData, oldFullLayout) {
@@ -32,5 +39,99 @@ exports.clean = function(newFullData, newFullLayout, oldFullData, oldFullLayout)
 
     if(hadPlot && !hasPlot) {
         oldFullLayout._paperdiv.selectAll('.sankey').remove();
+        oldFullLayout._paperdiv.selectAll('.bgsankey').remove();
     }
 };
+
+exports.updateFx = function(gd) {
+    for(var i = 0; i < gd._fullData.length; i++) {
+        subplotUpdateFx(gd, i);
+    }
+};
+
+function subplotUpdateFx(gd, index) {
+    var trace = gd._fullData[index];
+    var fullLayout = gd._fullLayout;
+
+    var dragMode = fullLayout.dragmode;
+    var cursor = fullLayout.dragmode === 'pan' ? 'move' : 'crosshair';
+    var bgRect = trace._bgRect;
+
+    if(dragMode === 'pan' || dragMode === 'zoom') return;
+
+    setCursor(bgRect, cursor);
+
+    var xaxis = {
+        _id: 'x',
+        c2p: Lib.identity,
+        _offset: trace._sankey.translateX,
+        _length: trace._sankey.width
+    };
+    var yaxis = {
+        _id: 'y',
+        c2p: Lib.identity,
+        _offset: trace._sankey.translateY,
+        _length: trace._sankey.height
+    };
+
+    // Note: dragOptions is needed to be declared for all dragmodes because
+    // it's the object that holds persistent selection state.
+    var dragOptions = {
+        gd: gd,
+        element: bgRect.node(),
+        plotinfo: {
+            id: index,
+            xaxis: xaxis,
+            yaxis: yaxis,
+            fillRangeItems: Lib.noop
+        },
+        subplot: index,
+        // create mock x/y axes for hover routine
+        xaxes: [xaxis],
+        yaxes: [yaxis],
+        doneFnCompleted: function(selection) {
+            var traceNow = gd._fullData[index];
+            var newGroups;
+            var oldGroups = traceNow.node.groups.slice();
+            var newGroup = [];
+
+            function findNode(pt) {
+                var nodes = traceNow._sankey.graph.nodes;
+                for(var i = 0; i < nodes.length; i++) {
+                    if(nodes[i].pointNumber === pt) return nodes[i];
+                }
+            }
+
+            for(var j = 0; j < selection.length; j++) {
+                var node = findNode(selection[j].pointNumber);
+                if(!node) continue;
+
+                // If the node represents a group
+                if(node.group) {
+                    // Add all its children to the current selection
+                    for(var k = 0; k < node.childrenNodes.length; k++) {
+                        newGroup.push(node.childrenNodes[k].pointNumber);
+                    }
+                    // Flag group for removal from existing list of groups
+                    oldGroups[node.pointNumber - traceNow.node._count] = false;
+                } else {
+                    newGroup.push(node.pointNumber);
+                }
+            }
+
+            newGroups = oldGroups
+                .filter(Boolean)
+                .concat([newGroup]);
+
+            Registry.call('_guiRestyle', gd, {
+                'node.groups': [ newGroups ]
+            }, index);
+        }
+    };
+
+    dragOptions.prepFn = function(e, startX, startY) {
+        prepSelect(e, startX, startY, dragOptions, dragMode);
+    };
+
+    dragElement.init(dragOptions);
+}

src/traces/sankey/calc.js

@@ -40,6 +40,7 @@ function convertToD3Sankey(trace) {
         if(linkSpec.target[i] > maxNodeId) maxNodeId = linkSpec.target[i];
     }
     var nodeCount = maxNodeId + 1;
+    trace.node._count = nodeCount;
 
     // Group nodes
     var j;

src/traces/sankey/index.js

@@ -18,6 +18,7 @@ Plot.plot = require('./plot');
 Plot.moduleType = 'trace';
 Plot.name = 'sankey';
 Plot.basePlotModule = require('./base_plot');
+Plot.selectPoints = require('./select.js');
 Plot.categories = ['noOpacity'];
 Plot.meta = {
     description: [

src/traces/sankey/render.js

@@ -836,6 +836,25 @@ module.exports = function(gd, svg, calcData, layout, callbacks) {
         .style('pointer-events', 'auto')
         .attr('transform', sankeyTransform);
 
+    sankey.each(function(d, i) {
+        gd._fullData[i]._sankey = d;
+        // Create dragbox if missing
+        var dragboxClassName = 'bgsankey-' + d.trace.uid + '-' + i;
+        Lib.ensureSingle(gd._fullLayout._draggers, 'rect', dragboxClassName);
+
+        gd._fullData[i]._bgRect = d3.select('.' + dragboxClassName);
+
+        // Style dragbox
+        gd._fullData[i]._bgRect
+          .style('pointer-events', 'all')
+          .attr('width', d.width)
+          .attr('height', d.height)
+          .attr('x', d.translateX)
+          .attr('y', d.translateY)
+          .classed('bgsankey', true)
+          .style({fill: 'transparent', 'stroke-width': 0});
+    });
+
     sankey.transition()
         .ease(c.ease).duration(c.duration)
         .attr('transform', sankeyTransform);
@@ -925,7 +944,8 @@ module.exports = function(gd, svg, calcData, layout, callbacks) {
         .call(attachPointerEvents, sankey, callbacks.nodeEvents)
         .call(attachDragHandler, sankeyLink, callbacks, gd); // has to be here as it binds sankeyLink
 
-    sankeyNode.transition()
+    sankeyNode
+        .transition()
         .ease(c.ease).duration(c.duration)
         .call(updateNodePositions)
         .style('opacity', function(n) { return n.partOfGroup ? 0 : 1;});

src/traces/sankey/select.js

@@ -0,0 +1,36 @@
+/**
+* Copyright 2012-2019, Plotly, Inc.
+* All rights reserved.
+*
+* This source code is licensed under the MIT license found in the
+* LICENSE file in the root directory of this source tree.
+*/
+
+'use strict';
+
+module.exports = function selectPoints(searchInfo, selectionTester) {
+    var cd = searchInfo.cd;
+    var selection = [];
+    var fullData = cd[0].trace;
+
+    var nodes = fullData._sankey.graph.nodes;
+
+    for(var i = 0; i < nodes.length; i++) {
+        var node = nodes[i];
+        if(node.partOfGroup) continue; // Those are invisible
+
+        // Position of node's centroid
+        var pos = [(node.x0 + node.x1) / 2, (node.y0 + node.y1) / 2];
+
+        // Swap x and y if trace is vertical
+        if(fullData.orientation === 'v') pos.reverse();
+
+        if(selectionTester && selectionTester.contains(pos, false, i, searchInfo)) {
+            selection.push({
+                pointNumber: node.pointNumber
+                // TODO: add eventData
+            });
+        }
+    }
+    return selection;
+};

test/jasmine/tests/sankey_test.js

@@ -379,6 +379,20 @@ describe('sankey tests', function() {
                 });
         });
 
+        it('Plotly.deleteTraces removes draggers', function(done) {
+            var mockCopy = Lib.extendDeep({}, mock);
+            Plotly.plot(gd, mockCopy)
+                .then(function() {
+                    expect(document.getElementsByClassName('bgsankey').length).toBe(1);
+                    return Plotly.deleteTraces(gd, [0]);
+                })
+                .then(function() {
+                    expect(document.getElementsByClassName('bgsankey').length).toBe(0);
+                })
+                .catch(failTest)
+                .then(done);
+        });
+
         it('Plotly.plot does not show Sankey if \'visible\' is false', function(done) {
             var mockCopy = Lib.extendDeep({}, mock);
 

test/jasmine/tests/select_test.js

@@ -13,6 +13,8 @@ var mouseEvent = require('../assets/mouse_event');
 var touchEvent = require('../assets/touch_event');
 
 var LONG_TIMEOUT_INTERVAL = 5 * jasmine.DEFAULT_TIMEOUT_INTERVAL;
+var delay = require('../assets/delay');
+var sankeyConstants = require('@src/traces/sankey/constants');
 
 function drag(path, options) {
     var len = path.length;
@@ -2666,6 +2668,64 @@ describe('Test select box and lasso per trace:', function() {
         .catch(failTest)
         .then(done);
     });
+
+    describe('should work on sankey traces', function() {
+        var waitingTime = sankeyConstants.duration * 2;
+
+        it('@flaky select', function(done) {
+            var fig = Lib.extendDeep({}, require('@mocks/sankey_circular.json'));
+            fig.layout.dragmode = 'select';
+            var dblClickPos = [250, 400];
+
+            Plotly.plot(gd, fig)
+            .then(function() {
+                // No groups initially
+                expect(gd._fullData[0].node.groups).toEqual([]);
+            })
+            .then(function() {
+                // Grouping the two nodes on the top right
+                return _run(
+                    [[640, 130], [400, 450]],
+                    function() {
+                        expect(gd._fullData[0].node.groups).toEqual([[2, 3]], 'failed to group #2 + #3');
+                    },
+                    dblClickPos, BOXEVENTS, 'for top right nodes #2 and #3'
+                );
+            })
+            .then(delay(waitingTime))
+            .then(function() {
+                // Grouping node #4 and the previous group
+                drag([[715, 400], [300, 110]]);
+            })
+            .then(delay(waitingTime))
+            .then(function() {
+                expect(gd._fullData[0].node.groups).toEqual([[4, 3, 2]], 'failed to group #4 + existing group of #2 and #3');
+            })
+            .catch(failTest)
+            .then(done);
+        });
+
+        it('@flaky should not work when dragmode is undefined', function(done) {
+            var fig = Lib.extendDeep({}, require('@mocks/sankey_circular.json'));
+            fig.layout.dragmode = undefined;
+
+            Plotly.plot(gd, fig)
+            .then(function() {
+                // No groups initially
+                expect(gd._fullData[0].node.groups).toEqual([]);
+            })
+            .then(function() {
+                // Grouping the two nodes on the top right
+                drag([[640, 130], [400, 450]]);
+            })
+            .then(delay(waitingTime))
+            .then(function() {
+                expect(gd._fullData[0].node.groups).toEqual([]);
+            })
+            .catch(failTest)
+            .then(done);
+        });
+    });
 });
 
 describe('Test that selections persist:', function() {