Allow setNodes to take a function for the label argument

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

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/lib/base-graph.js b/lib/base-graph.js
@@ -148,7 +148,11 @@ BaseGraph.prototype.setNodes = function(vs, label) {
   var args = arguments;
   _.each(vs, function(v) {
     if (args.length > 1) {
-      this.setNode(v, label);
+      if (_.isFunction(label)) {
+        this.setNode(v, label(v));
+      } else {
+        this.setNode(v, label);
+      }
     } else {
       this.setNode(v);
     }

diff --git a/test/base-graph-test.js b/test/base-graph-test.js
@@ -132,6 +132,15 @@ exports.tests = function(GraphConstructor) {
         expect(g.getNode("c")).to.be.undefined;
       });
 
+      it("can take a function for the label argument", function() {
+        g.setNodes(["a", "b", "c"], function(v) {
+          return v + "-label";
+        });
+        expect(g.getNode("a")).to.equal("a-label");
+        expect(g.getNode("b")).to.equal("b-label");
+        expect(g.getNode("c")).to.equal("c-label");
+      });
+
       it("is chainable", function() {
         var g2 = g.setNode("key", "label");
         expect(g).to.equal(g2);