1 files changed, 50 insertions, 9 deletions

diff --git a/lib/etatsbasen.js b/lib/etatsbasen.js
index bdd1dd9..cb29076 100644
--- a/lib/etatsbasen.js
+++ b/lib/etatsbasen.js
@@ -15,14 +15,19 @@ function emailIsInvalid(email) {
   return !re.test(email);
 }
 
+/** Filter the CSV data
+ * @param {object} csvdata - The CSV data
+ * @param  {number[]} - inclueOrgstructIds - Only include organizations with these organization structure ids
+ * @param {(number[]|string[]} headerFilter - Only include these headers
+ */
 function filter(csvdata, includeOrgstructIds) {
   if (! csvdata) {
     throw new TypeError('Missing `data` argument');
   }
   var headernameToIndex = {};
 
-  if (includeOrgstructIds && !Array.isArray(includeOrgstructIds)) {
-    throw new TypeError('`includedOrgstructIds` should be an array');
+  if (!includeOrgstructIds || !Array.isArray(includeOrgstructIds)) {
+    includeOrgstructIds = [];
   }
 
   csvdata = csvdata.filter(function(item, i) {
@@ -85,9 +90,14 @@ function renameHeader(data) {
   return data;
 }
 
-function removeColumns(data) {
-  if (! data) {
-    throw new TypeError('Missing `data` argument');
+function removeColumns(csvdata, headerFilter) {
+  var headernameToIndex = {};
+  if (! csvdata) {
+    throw new TypeError('Missing `csvdata` argument');
+  }
+
+  if (!headerFilter || !Array.isArray(headerFilter)) {
+    headerFilter = [];
   }
 
   var columns = ['url_nb', 'url_en', 'kommunenummer',
@@ -95,9 +105,13 @@ function removeColumns(data) {
 
   columns.forEach(function(name) {
     var toRemove = -1;
-    data.forEach(function(item, i) {
+    csvdata.forEach(function(item, i) {
       if (0 === i) {
         toRemove = item.indexOf(name);
+        // take note of headernames
+        item.forEach(function(item, i) {
+          headernameToIndex[item] = i;
+        });
       }
 
       if (toRemove >= 0) {
@@ -106,7 +120,31 @@ function removeColumns(data) {
     });
   });
 
-  return data;
+  if (headerFilter.length) {
+    // expand names to number
+    headerFilter = headerFilter.map(function(item) {
+      if (item) {
+        if (null !== item.toString().match(/^\d+$/)) {
+          return Number(item);
+        } else if (headernameToIndex[item]) {
+          return headernameToIndex[item];
+        } else {
+          throw 'Can\'t find header \'' + item + '\'. Possible values: ' + Object.keys(headernameToIndex).join(',');
+        }
+      }
+    });
+
+    // filter csvdata
+    csvdata = csvdata.map(function(item, i) {
+      var ret = [];
+      for (i = 0; i < headerFilter.length; i++) {
+        ret.unshift(item[headerFilter[i]]);
+      }
+      return ret;
+    });
+  }
+
+  return csvdata;
 }
 
 function addTags(data) {
@@ -177,16 +215,19 @@ exports.printCSV = function(cb, options) {
     return false;
   }
 
+  // This is just for fun
   csv().from.path(filename, { comment: '#'}).to.array( function(data) {
     print(
       removeColumns(
         addURL(
           addTags(
             renameHeader(
-              filter(data, options.categories)
+              filter(data,
+                     options.categories
+                    )
             )
           )
-        )
+        ), options.headers
       )
     );
     cb();