Add step labeling to sankey (spotify#73)

* Add step labeling functionality to sankey

* Update tests for sankey to include step labels

* Update docs

* Add step labeling functionality to sankey

* Update tests for sankey to include step labels

* Update docs

* PR review + documentation

* Better ternary

docs/build/index.html

@@ -31,5 +31,5 @@
     </div>
 
     <div class="container-fluid" id="container">Loading...</div>
-<script type="text/javascript" src="bundle.9ad33569594b4bef01a9.js"></script></body>
+<script type="text/javascript" src="bundle.46b96244f6acacb77226.js"></script></body>
 </html>

docs/src/docs/SankeyDiagram/examples/SankeyInteractive.js.example

@@ -13,10 +13,15 @@ class SankeyInteractiveExample extends React.Component {
         <SankeyDiagram
           width={900}
           height={500}
+          marginTop={50}
+          marginRight={50}
+          stepLabelText={step =>  `Step: ${step}`}
+          stepLabelPadding={16}
           nodes={graph.nodes}
           links={graph.links}
           nodeId={getNodeId}
           nodeAlignment="left"
+          nodeLabelPlacement="after"
           nodeLabelText={getNodeLabel}
           nodeStyle={(node) => {
             const nodeId = getNodeId(node);

docs/src/docs/SankeyDiagram/propDocs.json

@@ -1027,6 +1027,62 @@
         "value": "\"98%\"",
         "computed": false
       }
+    },
+    "stepLabelText": {
+      "type": {
+        "name": "union",
+        "value": [
+          {
+            "name": "string"
+          },
+          {
+            "name": "func"
+          }
+        ]
+      },
+      "required": false,
+      "description": "Text for step label or\naccessor function `f(step)` that returns the label text"
+    },
+    "stepLabelClassName": {
+      "type": {
+        "name": "union",
+        "value": [
+          {
+            "name": "string"
+          },
+          {
+            "name": "func"
+          }
+        ]
+      },
+      "required": false,
+      "description": "`className` attribute applied to each label,\nor accessor function which returns a class (string)"
+    },
+    "stepLabelStyle": {
+      "type": {
+        "name": "union",
+        "value": [
+          {
+            "name": "object"
+          },
+          {
+            "name": "func"
+          }
+        ]
+      },
+      "required": false,
+      "description": "Inline style object to be applied to each label,\nor accessor function which returns an object"
+    },
+    "stepLabelPadding": {
+      "type": {
+        "name": "number"
+      },
+      "required": false,
+      "description": "Vertical padding (in pixels) between step label and uppermost positioned node of that step",
+      "defaultValue": {
+        "value": "8",
+        "computed": false
+      }
     }
   }
 }

src/SankeyDiagram.js

@@ -210,6 +210,36 @@ const SankeyLinkLabel = props => {
   );
 };
 
+const SankeyStepLabel = props => {
+  const {
+    x,
+    y,
+    stepLabelPadding,
+    stepLabelText,
+    stepLabelClassName,
+    stepLabelStyle,
+    step
+  } = props;
+
+  let yPos = y;
+
+  if (_.isNumber(stepLabelPadding)) {
+    yPos = yPos - stepLabelPadding;
+  }
+
+  return (
+    <text
+      className={`rct-step-label ${getValue(stepLabelClassName, step)}`}
+      style={getValue(stepLabelStyle, step)}
+      x={x}
+      y={yPos}
+      key={`step-${x}-${step}`}
+    >
+      {getValue(stepLabelText, step)}
+    </text>
+  );
+};
+
 const SVGContainer = props => {
   const otherProps = _.omit(props, ["standalone"]);
   if (props.standalone) {
@@ -649,7 +679,26 @@ export default class SankeyDiagram extends React.Component {
     linkTargetLabelStartOffset: PropTypes.oneOfType([
       PropTypes.string,
       PropTypes.number
-    ])
+    ]),
+    /**
+     * Text for step label or
+     * accessor function `f(step)` that returns the label text
+     */
+    stepLabelText: PropTypes.oneOfType([PropTypes.string, PropTypes.func]),
+    /**
+     * `className` attribute applied to each label,
+     * or accessor function which returns a class (string)
+     */
+    stepLabelClassName: PropTypes.oneOfType([PropTypes.string, PropTypes.func]),
+    /**
+     * Inline style object to be applied to each label,
+     * or accessor function which returns an object
+     */
+    stepLabelStyle: PropTypes.oneOfType([PropTypes.object, PropTypes.func]),
+    /**
+     * Vertical padding (in pixels) between step label and uppermost positioned node of that step
+     */
+    stepLabelPadding: PropTypes.number
 
     //standalone
   };
@@ -668,6 +717,7 @@ export default class SankeyDiagram extends React.Component {
     showNodes: true,
     nodeWidth: 12,
     nodePadding: 8,
+    stepLabelPadding: 8,
     nodeAlignment: "justify",
     nodeClassName: "",
     nodeStyle: {},
@@ -849,13 +899,64 @@ export default class SankeyDiagram extends React.Component {
       );
     }
 
+    function displayStepLabelsIf(
+      stepLabelText,
+      stepLabelClassName,
+      stepLabelStyle,
+      stepLabelPadding,
+      nodes
+    ) {
+      if (!stepLabelText) {
+        return null;
+      }
+
+      const depthMapXPos = {};
+      const depthMapYPos = {};
+
+      nodes.forEach(n => {
+        depthMapXPos[n.depth] = n.x0;
+
+        // For the given depth, set the y equal to the highest positioned y value
+        depthMapYPos[n.depth] = depthMapYPos[n.depth]
+          ? Math.min(n.y0, depthMapYPos[n.depth])
+          : n.y0;
+      });
+
+      return (
+        <g className="rct-step-labels" width={innerWidth} height={100}>
+          {_.map(depthMapXPos, (x, step) => {
+            const stepLabelProps = {
+              y: depthMapYPos[step],
+              step,
+              x,
+              stepLabelText,
+              stepLabelClassName,
+              stepLabelPadding,
+              stepLabelStyle
+            };
+
+            return (
+              <SankeyStepLabel key={`rct-step-${step}`} {...stepLabelProps} />
+            );
+          })}
+        </g>
+      );
+    }
+
     return (
       <SVGContainer {...{ standalone, width, height, className, style }}>
         <g
           width={innerWidth}
           height={innerHeight}
           transform={`translate(${marginLeft}, ${marginTop})`}
         >
+          {displayStepLabelsIf(
+            this.props.stepLabelText,
+            this.props.stepLabelClassName,
+            this.props.stepLabelStyle,
+            this.props.stepLabelPadding,
+            graph.nodes
+          )}
           {mapLinksInGroupIf(
             this.props.showLinks,
             "rct-sankey-links",

tests/jsdom/spec/SankeyDiagram.spec.js

@@ -17,6 +17,7 @@ const SankeyLink = Sankey.__get__("SankeyLink");
 const SankeyNodeTerminal = Sankey.__get__("SankeyLink");
 const SankeyNodeLabel = Sankey.__get__("SankeyNodeLabel");
 const SankeyLinkLabel = Sankey.__get__("SankeyLinkLabel");
+const SankeyStepLabel = Sankey.__get__("SankeyStepLabel");
 
 function getSampleData() {
   return {
@@ -393,6 +394,30 @@ describe("SankeyDiagram", () => {
     });
   });
 
+  it("passes stepLabelText, stepLabelClassName, stepLabelPadding and stepLabelStyle through to SankeyStepLabel", () => {
+    const props = {
+      ...getSampleData(),
+      width: 600,
+      height: 400,
+      stepLabelText: "text",
+      stepLabelClassName: "scoop",
+      stepLabelStyle: { fill: "orange" },
+      stepLabelPadding: 16
+    };
+    const chart = mount(<SankeyDiagram {...props} />);
+    const sankeyStepLabels = chart.find(SankeyStepLabel);
+    expect(sankeyStepLabels).to.have.length(3);
+
+    sankeyStepLabels.forEach((label, i) => {
+      const stepLabelProps = label.props();
+      expect(stepLabelProps.stepLabelText).to.equal("text");
+      expect(stepLabelProps.stepLabelClassName).to.equal("scoop");
+      expect(stepLabelProps.stepLabelPadding).to.equal(16);
+      expect(stepLabelProps.stepLabelStyle).to.be.an("object");
+      expect(stepLabelProps.stepLabelStyle.fill).to.equal("orange");
+    });
+  });
+
   it("uses showNodeLabels boolean or accessor prop to determine whether to render node labels", () => {
     const size = { width: 600, height: 400 };
     const showNodeLabelsProps = {
@@ -459,6 +484,17 @@ describe("SankeyDiagram", () => {
     expect(showSomeLinkLabelsChart.find(SankeyLinkLabel)).to.have.length(2);
   });
 
+  it("uses stepLabelText text or accessor prop to determine whether to render SankeyStepLabels", () => {
+    const size = { width: 600, height: 400 };
+    const stepLabelsProps = {
+      ...size,
+      ...getSampleData(),
+      stepLabelText: step => `Step: ${step}`
+    };
+    const stepLabelsChart = mount(<SankeyDiagram {...stepLabelsProps} />);
+    expect(stepLabelsChart.find(SankeyStepLabel)).to.have.length(3);
+  });
+
   describe("SankeyNode", () => {
     const basicNodeObj = {
       index: 5,
@@ -929,6 +965,68 @@ describe("SankeyDiagram", () => {
       expect(textPath.props().startOffset).to.equal("27%");
     });
   });
+
+  describe("SankeyStepLabel", () => {
+    const step = 0;
+
+    it("renders a step label", () => {
+      const props = {
+        step,
+        x: 100,
+        y: 100,
+        stepLabelText: () => "r2d2"
+      };
+      const label = mount(<SankeyStepLabel {...props} />);
+
+      const text = label.find("text");
+      expect(text).to.have.length(1);
+      expect(text.props().x).to.be.finite;
+      expect(text.props().y).to.be.finite;
+      expect(text.text()).to.equal("r2d2");
+    });
+    it("uses stepLabelText accessor prop to create label text", () => {
+      const labelWithName = mount(
+        <SankeyStepLabel
+          {...{
+            step,
+            stepLabelText: step => step
+          }}
+        />
+      );
+      const textWithName = labelWithName.find("text");
+      expect(textWithName).to.have.length(1);
+      expect(textWithName.text()).to.equal("0");
+    });
+    it("passes stepLabelClassName & stepLabelStyle through to the text element", () => {
+      const props = {
+        step,
+        stepLabelText: () => "r2d2",
+        stepLabelClassName: "link-zelda",
+        stepLabelStyle: { fill: "orange" }
+      };
+      const label = mount(<SankeyStepLabel {...props} />);
+      const text = label.find("text");
+      expect(text).to.have.length(1);
+      expect(text.props().className).to.contain("link-zelda");
+      expect(text.props().style).to.be.an("object");
+      expect(text.props().style.fill).to.equal("orange");
+    });
+    it("calls stepLabelClassName & stepLabelStyle if they are functions", () => {
+      const props = {
+        step,
+        stepLabelText: () => "r2d2",
+        stepLabelClassName: step => `step-${step}`,
+        stepLabelStyle: () => ({ fill: "thistle" })
+      };
+      const label = mount(<SankeyStepLabel {...props} />);
+      const text = label.find("text");
+      expect(text).to.have.length(1);
+      expect(text.props().className).to.contain("step-0");
+      expect(text.props().style).to.be.an("object");
+      expect(text.props().style.fill).to.equal("thistle");
+    });
+  });
+
   // todo test terminals
   // test their properties & rendered correctly
 });