3 files changed, 95 insertions, 17 deletions

diff --git a/bin/etatsbasen.js b/bin/etatsbasen.js
index 9241e10..0960900 100755
--- a/bin/etatsbasen.js
+++ b/bin/etatsbasen.js
@@ -7,6 +7,8 @@ var pkg = require(path.join(__dirname, '..', 'package.json'));
 var argv = require('minimist')(process.argv.slice(2));
 var etatsbasen = require('etatsbasen');
 
+var defaultCategories = [12,14,17,18,27,33,38,66,68,76];
+
 process.bin = process.title = 'etatsbasen';
 
 if (argv.v) {
@@ -17,10 +19,28 @@ if (argv.v) {
 var options = {};
 options.filename = argv.f || 'etatsbasen.csv';
 
-if (argv.c && 'string' === typeof argv.c) {
-  options.categories = argv.c.split(',');
-} else {
-  options.categories = [];
+if (argv.c) {
+  if (Array.isArray(argv.c)) {
+    options.categories = argv.c;
+ } else if ('string' === typeof argv.c) {
+   if ('all' !== argv.c) {
+     options.categories = [argv.c];
+   }
+ } else {
+   options.categories = defaultCategories;
+ }
+}
+
+if (argv.o) {
+  if (Array.isArray(argv.o)) {
+    options.headers = argv.o;
+  } else if ('string' === typeof argv.o ||
+            'number' === typeof argv.o) {
+    options.headers = [argv.o];
+  } else {
+    console.log('Unable to parse -o option(s)');
+    process.exit(1);
+  }
 }
 
 function fileNotFound() {
@@ -35,11 +55,13 @@ if (argv.h || fileNotFound()) {
   console.log([
     'usage: etatsbasen [options]',
     '',
-    '  -c [c1,c2,..]   Categories to include',
-    '  -f [file]       File to read from (defaults: `etatsbasen.csv`)',
-    '  -v              Print version.',
-    '  -h              Write this help.'
+    '  -c [all|[c1,c2,..]]   Categories to include (defaults: `' + defaultCategories.join(',') + '`)',
+    '  -f [file]             File to read from (defaults: `etatsbasen.csv`)',
+    '  -o h1[,h2,h3]         Include only these headers in output (id or name)',
+    '  -v                    Print version.',
+    '  -h                    Write this help.'
   ].join('\n'));
+
   process.exit();
 }
 
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();
diff --git a/test/etatsbasen_test.js b/test/etatsbasen_test.js
index 1b7c7da..6f42db7 100644
--- a/test/etatsbasen_test.js
+++ b/test/etatsbasen_test.js
@@ -74,6 +74,21 @@ exports.etatsbasen = {
         test.ok(!err);
         test.done();
       }, { filename: 'fixtures/2.csv', categories: [36,37] });
+    },
+    'filter headers': function(test) {
+      var oldLogger = console.log;
+      console.log = function(str) {
+        test.expect(5);
+        test.ok(str.match(/.*request\_email.*/), 'Can\'t find filtered header email');
+        test.equal(str.split('\n').length, 10);
+        test.ok(!str.match(/\,\,/), 'Found empty entries (WARNING: suspect test)');
+        test.ok(str.match(/.*name\.nn.*/, 'Missing name.nn header'));
+      };
+      etatsbasen.printCSV(function(err) {
+        console.log = oldLogger;
+        test.ok(!err);
+        test.done();
+      }, { filename: 'fixtures/2.csv', headers: ['request_email', 3] });
     }
   }
   };