Add preExportItemFunction

TransitApp · Mar 9, 2018 · 0c2e98e · 0c2e98e
1 parent 6ca3161
commit 0c2e98e
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 11 deletions.
diff --git a/gtfs.js b/gtfs.js
@@ -185,14 +185,21 @@ class Gtfs {
    * @param {{
    *   regexPatternObjectsByTableName: Map.<string, Array.<{regex: RegExp, pattern: string}>>,
    *   throws: boolean,
-   *   postImportTableFunction: function,
    *   forcedSchema,
+   *   postImportItemFunction: function,
+   *   preExportItemFunction: function,
    * }} [options] Optional. See list above.
    * @return {Gtfs} gtfs Instanciated GTFS object.
    */
   constructor(
     path,
-    { regexPatternObjectsByTableName = new Map(), throws = true, postImportTableFunction, forcedSchema } = {}
+    {
+      regexPatternObjectsByTableName = new Map(),
+      throws = true,
+      forcedSchema,
+      postImportItemFunction,
+      preExportItemFunction,
+    } = {}
   ) {
     if (path !== undefined) {
       if (typeof path !== 'string' || path.length === 0) {
@@ -211,7 +218,8 @@ class Gtfs {
     this._path = path;
     this._regexPatternObjectsByTableName = regexPatternObjectsByTableName;
     this._shouldThrow = throws;
-    this._postImportTableFunction = postImportTableFunction;
+    this._postImportItemFunction = postImportItemFunction;
+    this._preExportItemFunction = preExportItemFunction;
     this._tables = new Map();
     this._schema = forcedSchema || defaultSchema;
   }

diff --git a/helpers/export.js b/helpers/export.js
@@ -82,7 +82,11 @@ function exportTable(tableName, gtfs, outputPath, callback) {
     const deepness = gtfs._schema.deepnessByTableName[tableName];
 
     if (deepness === 0) {
-      const row = fromObjectToCsvString(gtfs.getIndexedTable(tableName), keys);
+      let item = gtfs.getIndexedTable(tableName);
+      if (gtfs._preExportItemFunction) {
+        item = gtfs._preExportItemFunction(item);
+      }
+      const row = fromObjectToCsvString(item, keys);
       fs.appendFile(outputFullPath, row, callback);
       return;
     }
@@ -91,9 +95,15 @@ function exportTable(tableName, gtfs, outputPath, callback) {
 
     async.eachSeries(gtfs.getIndexedTable(tableName), acomb.ensureAsync(([key, object], subDone) => {
       if (deepness === 1) {
+        if (gtfs._preExportItemFunction) {
+          object = gtfs._preExportItemFunction(object);
+        }
         rowsBuffer.push(fromObjectToCsvString(object, keys));
       } else if (deepness === 2) {
         object.forEach((subObject) => {
+          if (gtfs._preExportItemFunction) {
+            subObject = gtfs._preExportItemFunction(subObject);
+          }
           rowsBuffer.push(fromObjectToCsvString(subObject, keys));
         });
       }

diff --git a/helpers/import.js b/helpers/import.js
@@ -30,8 +30,8 @@ exports.importTable = (gtfs, tableName) => {
     gtfs._tables.set(tableName, new Map());
   }
 
-  if (gtfs._postImportTableFunction) {
-    gtfs.forEachItemInTable(tableName, gtfs._postImportTableFunction);
+  if (gtfs._postImportItemFunction) {
+    gtfs.forEachItemInTable(tableName, gtfs._postImportItemFunction);
   }
 };
 
@@ -50,7 +50,7 @@ function getRows(buffer, regexPatternObjects, tableName) {
     rowsSlice = buffer.toString('utf8', position, Math.min(buffer.length, position + batchLength));
 
     if (regexPatternObjects) {
-      regexPatternObjects.forEach(({regex, pattern}) => {
+      regexPatternObjects.forEach(({ regex, pattern }) => {
         const modifiedRowsSlice = rowsSlice.replace(regex, pattern || '');
 
         if (modifiedRowsSlice !== rowsSlice) {

diff --git a/tests/gtfs_options.js b/tests/gtfs_options.js
@@ -4,6 +4,7 @@
 
 // eslint-disable-next-line import/no-extraneous-dependencies
 const { expect } = require('chai');
+const fs = require('fs-extra');
 
 const { Gtfs } = require('../index');
 
@@ -36,13 +37,68 @@ describe('Tests on GTFS constructor options', () => {
 
   it('Test on postImportTableFunction', (done) => {
     const path = `${__dirname}/samples/1/`;
-    const postImportTableFunction = (item) => { item.temp = 'some value'; };
-    const gtfs = new Gtfs(path, { postImportTableFunction });
+    const postImportItemFunction = (item) => { item.temp = 'some value'; };
+    const gtfs = new Gtfs(path, { postImportItemFunction });
 
-    const route = gtfs.getRouteWithId('route_0')
+    const route = gtfs.getRouteWithId('route_0');
 
     expect(route.temp).to.equal('some value');
 
-    done();
+    const outputPath = `${__dirname}/temp_4865de67d01696s48dfbd0e71adx8f0b/`;
+    gtfs.exportAtPath(outputPath, (exportError) => {
+      if (exportError) { throw exportError; }
+
+      fs.readFile(`${outputPath}routes.txt`, (readRoutesError, routesTxt) => {
+        if (readRoutesError) { throw readRoutesError; }
+
+        expect(String(routesTxt)).to.equal(
+          'route_id,agency_id,route_short_name,route_long_name,route_desc,route_type,route_url,route_color,' +
+          'route_text_color,route_sort_order,temp\n' +
+          'route_0,agency_0,R0,Route 0,,3,,,,,some value\n'
+        );
+
+        fs.remove(outputPath, (removeError) => {
+          if (removeError) { throw removeError; }
+
+          done();
+        });
+      });
+    });
+  });
+
+  it('Test on postImportTableFunction', (done) => {
+    const path = `${__dirname}/samples/1/`;
+    const postImportItemFunction = (item) => { item.temp = { key: 'value' }; };
+    const preExportItemFunction = (item) => {
+      const item2 = JSON.parse(JSON.stringify(item));
+      item2.temp = JSON.stringify(item.temp);
+      return item2;
+    };
+    const gtfs = new Gtfs(path, { postImportItemFunction, preExportItemFunction });
+
+    const route = gtfs.getRouteWithId('route_0');
+
+    expect(route.temp).to.deep.equal({ key: 'value' });
+
+    const outputPath = `${__dirname}/temp_4865de67d01f96s489fbd0e71ad88f0b/`;
+    gtfs.exportAtPath(outputPath, (exportError) => {
+      if (exportError) { throw exportError; }
+
+      fs.readFile(`${outputPath}routes.txt`, (readRoutesError, routesTxt) => {
+        if (readRoutesError) { throw readRoutesError; }
+
+        expect(String(routesTxt)).to.equal(
+          'route_id,agency_id,route_short_name,route_long_name,route_desc,route_type,route_url,route_color,' +
+          'route_text_color,route_sort_order,temp\n' +
+          'route_0,agency_0,R0,Route 0,,3,,,,,"{""key"":""value""}"\n'
+        );
+
+        fs.remove(outputPath, (removeError) => {
+          if (removeError) { throw removeError; }
+
+          done();
+        });
+      });
+    });
   });
 });