feat(sankey): improve support for nodes sorting

plouc · Mar 23, 2019 · f63450f · f63450f
1 parent fec85fa
commit f63450f
Show file tree

Hide file tree

Showing 9 changed files with 231 additions and 198 deletions.
diff --git a/packages/sankey/index.d.ts b/packages/sankey/index.d.ts
@@ -40,6 +40,8 @@ declare module '@nivo/sankey' {
         | 'color'
         | 'luminosity'
 
+    export type SankeySortFunction = (nodeA: SankeyDataNode, nodeB: SankeyDataNode) => number
+
     export type SankeyProps = Partial<{
         align: 'center' | 'justify' | 'left' | 'right'
 
@@ -77,7 +79,7 @@ declare module '@nivo/sankey' {
 
         legends: LegendProps[]
 
-        sort?: (nodeA: SankeyDataNode, nodeB: SankeyDataNode) => number
+        sort: 'auto' | 'input' | 'ascending' | 'descending' | SankeySortFunction
     }>
 
     interface Dimensions {

diff --git a/packages/sankey/package.json b/packages/sankey/package.json
@@ -25,7 +25,7 @@
   "dependencies": {
     "@nivo/core": "0.53.0",
     "@nivo/legends": "0.53.0",
-    "d3-sankey": "0.12.1",
+    "d3-sankey": "^0.12.1",
     "lodash": "^4.17.4",
     "react-motion": "^0.5.2",
     "recompose": "^0.30.0"

diff --git a/packages/sankey/src/Sankey.js b/packages/sankey/src/Sankey.js
@@ -24,6 +24,7 @@ const Sankey = ({
     data: _data,
 
     align,
+    sortFunction,
 
     margin,
     width,
@@ -74,12 +75,10 @@ const Sankey = ({
     tooltipFormat,
 
     legends,
-
-    sort,
 }) => {
     const sankey = d3Sankey()
-        .nodeSort(sort)
         .nodeAlign(sankeyAlignmentFromProp(align))
+        .nodeSort(sortFunction)
         .nodeWidth(nodeWidth)
         .nodePadding(nodePaddingY)
         .size([width, height])

diff --git a/packages/sankey/src/enhance.js b/packages/sankey/src/enhance.js
@@ -38,5 +38,17 @@ export default Component =>
         withPropsOnChange(['label', 'labelFormat'], ({ label, labelFormat }) => ({
             getLabel: getLabelGenerator(label, labelFormat),
         })),
+        withPropsOnChange(['sort'], ({ sort }) => {
+            let sortFunction
+            if (sort === 'input') {
+                sortFunction = null
+            } else if (sort === 'ascending') {
+                sortFunction = (a, b) => a.value - b.value
+            } else if (sort === 'descending') {
+                sortFunction = (a, b) => b.value - a.value
+            }
+
+            return { sortFunction }
+        }),
         pure
     )(Component)
diff --git a/packages/sankey/src/props.js b/packages/sankey/src/props.js
@@ -40,6 +40,10 @@ export const SankeyPropTypes = {
     }).isRequired,
 
     align: sankeyAlignmentPropType.isRequired,
+    sort: PropTypes.oneOfType([
+        PropTypes.oneOf(['auto', 'input', 'ascending', 'descending']),
+        PropTypes.func,
+    ]).isRequired,
 
     nodeOpacity: PropTypes.number.isRequired,
     nodeHoverOpacity: PropTypes.number.isRequired,
@@ -79,6 +83,7 @@ export const SankeyPropTypes = {
 
 export const SankeyDefaultProps = {
     align: 'center',
+    sort: 'auto',
 
     nodeOpacity: 0.75,
     nodeHoverOpacity: 1,