Provide mechanism to set default labels

kkowar · Sep 24, 2014 · 8ad1e8e · 8ad1e8e
1 parent 7e7d7ed
commit 8ad1e8e
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 4 deletions.
diff --git a/lib/base-graph.js b/lib/base-graph.js
@@ -17,6 +17,12 @@ function BaseGraph(inEdges, outEdges, sources, sinks) {
 
   // v -> true
   this._sinks = sinks;
+
+  // A function of v -> label that is used when creating a node with no label
+  this._defaultNodeLabelFn = _.constant(undefined);
+
+  // A function of {v, w} -> label that is used when creating an edge with no label
+  this._defaultEdgeLabelFn = _.constant(undefined);
 }
 
 /* Number of nodes in the graph. Should only be changed by the implementation. */
@@ -126,6 +132,20 @@ BaseGraph.prototype.edges = function() {
   }, new Array(this._edgeCount));
 };
 
+BaseGraph.prototype.setDefaultNodeLabel = function(label) {
+  if (!_.isFunction(label)) {
+    label = _.constant(label);
+  }
+  this._defaultNodeLabelFn = label;
+};
+
+BaseGraph.prototype.setDefaultEdgeLabel = function(label) {
+  if (!_.isFunction(label)) {
+    label = _.constant(label);
+  }
+  this._defaultEdgeLabelFn = label;
+};
+
 
 /* === Node-level functions ========== */
 
@@ -135,7 +155,7 @@ BaseGraph.prototype.setNode = function(v, label) {
     this._inEdges[v] = {};
     this._outEdges[v] = {};
     this._sources[v] = this._sinks[v] = true;
-    nodes[v] = label;
+    nodes[v] = arguments.length > 1 ? label : this._defaultNodeLabelFn(v);
     ++this._nodeCount;
     this._onAddNode(v);
   } else if (arguments.length > 1) {
@@ -228,15 +248,18 @@ BaseGraph.prototype.setEdge = function(v, w, label) {
   var outV = this._outEdges[v],
       prevCtx = outV[w];
   if (prevCtx === undefined) {
+    if (arguments.length < 3) {
+      label = this._defaultEdgeLabelFn({ v: v, w: w });
+    }
     ++this._edgeCount;
   }
 
   var edgeCtx = { v: String(v), w: String(w) };
-  if (arguments.length <= 2) {
-    if (_.has(prevCtx, "label")) {
+  if (label === undefined) {
+    if (arguments.length < 3 && _.has(prevCtx, "label")) {
       edgeCtx.label = prevCtx.label;
     }
-  } else if (label !== undefined) {
+  } else {
     edgeCtx.label = label;
   }
   Object.freeze(edgeCtx);

diff --git a/test/base-graph-test.js b/test/base-graph-test.js
@@ -178,6 +178,24 @@ exports.tests = function(GraphConstructor) {
       });
     });
 
+    describe("setDefaultNodeLabel", function() {
+      it("assigns the default value for created nodes", function() {
+        g.setDefaultNodeLabel(1);
+        g.setNode("a");
+        g.setNode("b");
+        expect(g.getNode("a")).to.equal(1);
+        expect(g.getNode("b")).to.equal(1);
+      });
+
+      it("can use a function to create the default value", function() {
+        g.setDefaultNodeLabel(function(v) { return v + "-label"; });
+        g.setNode("a");
+        g.setNode("b");
+        expect(g.getNode("a")).to.equal("a-label");
+        expect(g.getNode("b")).to.equal("b-label");
+      });
+    });
+
     describe("removeNode", function() {
       it("does nothing if the node is not part of the graph", function() {
         var removed;
@@ -434,6 +452,21 @@ exports.tests = function(GraphConstructor) {
       });
     });
 
+    describe("setDefaultEdgeLabel", function() {
+      it("assigns the default value for created edges", function() {
+        g.setDefaultEdgeLabel(1);
+        g.setEdge("a", "b");
+        expect(g.getEdge("a", "b")).to.equal(1);
+      });
+
+      it("can use a function to create the default value", function() {
+        g.setDefaultEdgeLabel(function(edge) {
+          return edge.v + "->" + edge.w;
+        });
+        g.setEdge("a", "b");
+        expect(g.getEdge("a", "b")).to.equal("a->b");
+      });
+    });
     describe("removeEdge", function() {
       it("does nothing if the edge is not in the graph", function() {
         g.setNode("n1");