Allow setPath to take a function for its label argument

kkowar · Sep 24, 2014 · fde8b40 · fde8b40
1 parent 53c93a0
commit fde8b40
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/lib/base-graph.js b/lib/base-graph.js
@@ -250,7 +250,11 @@ BaseGraph.prototype.setPath = function(vs, label) {
       args = arguments;
   _.reduce(vs, function(v, w) {
     if (args.length > 1) {
-      self.setEdge(v, w, label);
+      if (_.isFunction(label)) {
+        self.setEdge(v, w, label({ v: v, w: w }));
+      } else {
+        self.setEdge(v, w, label);
+      }
     } else {
       self.setEdge(v, w);
     }

diff --git a/test/base-graph-test.js b/test/base-graph-test.js
@@ -381,6 +381,12 @@ exports.tests = function(GraphConstructor) {
         expect(g.getEdge("b", "c")).to.be.undefined;
       });
 
+      it("can take a function as the label argument", function() {
+        g.setPath(["a", "b", "c"], function(edge) { return edge.v + "-" + edge.w; });
+        expect(g.getEdge("a", "b")).to.equal("a-b");
+        expect(g.getEdge("b", "c")).to.equal("b-c");
+      });
+
       it("is chainable", function() {
         var g2 = g.setPath(["a", "b", "c"]);
         expect(g).to.equal(g2);