Tree: let the width parameter also be a function that returns an int

1 files changed, 13 insertions, 2 deletions

diff --git a/lib/tree.js b/lib/tree.js
index 75e2d75f..839b698c 100644
--- a/lib/tree.js
+++ b/lib/tree.js
@@ -25,6 +25,7 @@ var tree = {
             };
         }
 
+
         if (width <= 0) {
             throw {
                 name: "TreeParamError",
@@ -46,9 +47,19 @@ var tree = {
                 else {
                     if (depth !== 0) {
                         value = [];
-                        for (var i = 0; i < width; i++) {
+                        var evalWidth = 1;
+
+                        if (typeof(width) == "function") {
+                            evalWidth = width();
+                        }
+                        else {
+                            evalWidth = width;
+                        }
+
+                        for (var i = 0; i < evalWidth; i++) {
                             value.push(this.createTree(depth - 1, width, obj));
                         }
+
                     }
                 }
 
@@ -61,4 +72,4 @@ var tree = {
 
 };
 
-module.exports = tree;
\ No newline at end of file
+module.exports = tree;