Add support for ES6 function notations (#16)

rhendric · blakeembrey · commit d67b657fe202 · 2019-02-04T20:40:54.000-08:00
diff --git a/javascript-stringify.js b/javascript-stringify.js
@@ -81,6 +81,38 @@
     return !RESERVED_WORDS.hasOwnProperty(name) && IS_VALID_IDENTIFIER.test(name);
   }
 
+  /**
+   * Check if a function is an ES6 generator function
+   *
+   * @param  {Function} fn
+   * @return {boolean}
+   */
+  function isGeneratorFunction (fn) {
+    return fn.constructor.name === 'GeneratorFunction';
+  }
+
+  /**
+   * Can be replaced with `str.startsWith(prefix)` if code is updated to ES6.
+   *
+   * @param  {string} str
+   * @param  {string} prefix
+   * @return {boolean}
+   */
+  function stringStartsWith (str, prefix) {
+    return str.substring(0, prefix.length) === prefix;
+  }
+
+  /**
+   * Can be replaced with `str.repeat(count)` if code is updated to ES6.
+   *
+   * @param  {string} str
+   * @param  {number} count
+   * @return {string}
+   */
+  function stringRepeat (str, count) {
+    return new Array(Math.max(0, count|0) + 1).join(str);
+  }
+
   /**
    * Return the global variable name.
    *
@@ -149,19 +181,59 @@
   function stringifyObject (object, indent, next) {
     // Iterate over object keys and concat string together.
     var values = Object.keys(object).reduce(function (values, key) {
-      var value = next(object[key], key);
+      var value;
+      var addKey = true;
+
+      // Handle functions specially to detect method notation.
+      if (typeof object[key] === 'function') {
+        var fn = object[key];
+        var fnString = fn.toString();
+        var prefix = isGeneratorFunction(fn) ? '*' : '';
+
+        // Was this function defined with method notation?
+        if (fn.name === key && stringStartsWith(fnString, prefix + key + '(')) {
+          if (isValidVariableName(key)) {
+            // The function is already in valid method notation.
+            value = fnString;
+          } else {
+            // Reformat the opening of the function into valid method notation.
+            value = prefix + stringify(key) + fnString.substring(prefix.length + key.length);
+          }
+
+          // Dedent the function, since it didn't come through regular stringification.
+          if (indent) {
+            value = dedentFunction(value);
+          }
 
-      // Omit `undefined` object values.
-      if (value === undefined) {
-        return values;
+          // Method notation includes the key, so there's no need to add it again below.
+          addKey = false;
+        } else {
+          // Not defined with method notation; delegate to regular stringification.
+          value = next(fn, key);
+        }
+      } else {
+        // `object[key]` is not a function.
+        value = next(object[key], key);
+
+        // Omit `undefined` object values.
+        if (value === undefined) {
+          return values;
+        }
       }
 
-      // String format the key and value data.
-      key   = isValidVariableName(key) ? key : stringify(key);
+      // String format the value data.
       value = String(value).split('\n').join('\n' + indent);
 
-      // Push the current object key and value into the values array.
-      values.push(indent + key + ':' + (indent ? ' ' : '') + value);
+      if (addKey) {
+        // String format the key data.
+        key = isValidVariableName(key) ? key : stringify(key);
+
+        // Push the current object key and value into the values array.
+        values.push(indent + key + ':' + (indent ? ' ' : '') + value);
+      } else {
+        // Push just the value; this is a method and no key is needed.
+        values.push(indent + value);
+      }
 
       return values;
     }, []).join(indent ? ',\n' : ',');
@@ -174,6 +246,50 @@
     return '{' + values + '}';
   }
 
+  /**
+   * Rewrite a stringified function to remove initial indentation.
+   *
+   * @param  {string} fnString
+   * @return {string}
+   */
+  function dedentFunction (fnString) {
+    var indentationRegExp = /\n */g;
+    var match;
+
+    // Find the minimum amount of indentation used in the function body.
+    var dedent = Infinity;
+    while (match = indentationRegExp.exec(fnString)) {
+      dedent = Math.min(dedent, match[0].length - 1);
+    }
+
+    if (isFinite(dedent)) {
+      return fnString.split('\n' + stringRepeat(' ', dedent)).join('\n');
+    } else {
+      // Function is a one-liner and needs no adjustment.
+      return fnString;
+    }
+  }
+
+  /**
+   * Stringify a function.
+   *
+   * @param  {Function} fn
+   * @return {string}
+   */
+  function stringifyFunction (fn, indent) {
+    var value = fn.toString();
+    if (indent) {
+      value = dedentFunction(value);
+    }
+    var prefix = isGeneratorFunction(fn) ? '*' : '';
+    if (fn.name && stringStartsWith(value, prefix + fn.name + '(')) {
+      // Method notation was used to define this function, but it was transplanted from another object.
+      // Convert to regular function notation.
+      value = 'function' + prefix + ' ' + value.substring(prefix.length);
+    }
+    return value;
+  }
+
   /**
    * Convert JavaScript objects into strings.
    */
@@ -202,7 +318,8 @@
       return 'new Map(' + stringify(Array.from(array), indent, next) + ')';
     },
     '[object RegExp]': String,
-    '[object Function]': String,
+    '[object Function]': stringifyFunction,
+    '[object GeneratorFunction]': stringifyFunction,
     '[object global]': toGlobalVariable,
     '[object Window]': toGlobalVariable
   };
@@ -263,7 +380,7 @@
 
     // Convert the spaces into a string.
     if (typeof space !== 'string') {
-      space = new Array(Math.max(0, space|0) + 1).join(' ');
+      space = stringRepeat(' ', space);
     }
 
     var maxDepth = Number(options.maxDepth) || 100;
diff --git a/package.json b/package.json
@@ -9,7 +9,7 @@
     "javascript-stringify.d.ts"
   ],
   "scripts": {
-    "test": "istanbul cover node_modules/mocha/bin/_mocha -- -R spec"
+    "test": "istanbul cover node_modules/mocha/bin/_mocha -x test.js -- -R spec"
   },
   "repository": "https://github.com/blakeembrey/javascript-stringify.git",
   "keywords": [
diff --git a/test.js b/test.js
@@ -9,6 +9,10 @@ describe('javascript-stringify', function () {
     };
   };
 
+  var testRoundTrip = function (insult, indent, options) {
+    return test(eval('(' + insult + ')'), insult, indent, options);
+  };
+
   describe('types', function () {
     describe('booleans', function () {
       it('should be stringified', test(true, 'true'));
@@ -82,6 +86,56 @@ describe('javascript-stringify', function () {
       );
     });
 
+    describe('functions', function () {
+      it(
+        'should reindent function bodies',
+        test(
+          function () {
+            if (true) {
+              return "hello";
+            }
+          },
+          'function () {\n  if (true) {\n    return "hello";\n  }\n}',
+          2
+        )
+      );
+
+      it(
+        'should reindent function bodies in objects',
+        test(
+          {
+            fn: function () {
+              if (true) {
+                return "hello";
+              }
+            }
+          },
+          '{\n  fn: function () {\n    if (true) {\n      return "hello";\n    }\n  }\n}',
+          2
+        )
+      );
+
+      it(
+        'should reindent function bodies in arrays',
+        test(
+          [
+            function () {
+              if (true) {
+                return "hello";
+              }
+            }
+          ],
+          '[\n  function () {\n    if (true) {\n      return "hello";\n    }\n  }\n]',
+          2
+        )
+      );
+
+      it(
+        'should not need to reindent one-liners',
+        testRoundTrip('{\n  fn: function () { return; }\n}', 2)
+      );
+    });
+
     describe('native instances', function () {
       describe('Date', function () {
         var date = new Date();
@@ -130,6 +184,103 @@ describe('javascript-stringify', function () {
             it('should stringify', test(new Set(['key', 'value']), "new Set(['key','value'])"));
           });
         }
+
+        describe('arrow functions', function () {
+          it('should stringify', testRoundTrip('(a, b) => a + b'));
+
+          it(
+            'should reindent function bodies',
+            test(
+              eval(
+'               (() => {\n' +
+'                 if (true) {\n' +
+'                   return "hello";\n' +
+'                 }\n' +
+'               })'),
+              '() => {\n  if (true) {\n    return "hello";\n  }\n}',
+              2
+            )
+          );
+        });
+
+        describe('generators', function () {
+          it('should stringify', testRoundTrip('function* (x) { yield x; }'));
+        });
+
+        describe('method notation', function () {
+          it('should stringify', testRoundTrip('{a(b, c) { return b + c; }}'));
+
+          it('should stringify generator methods', testRoundTrip('{*a(b) { yield b; }}'));
+
+          it(
+            'should not be fooled by tricky names',
+            testRoundTrip("{'function a'(b, c) { return b + c; }}")
+          );
+
+          it(
+            'should not be fooled by tricky generator names',
+            testRoundTrip("{*'function a'(b, c) { return b + c; }}")
+          );
+
+          it(
+            'should not be fooled by empty names',
+            testRoundTrip("{''(b, c) { return b + c; }}")
+          );
+
+          it(
+            'should not be fooled by arrow functions',
+            testRoundTrip("{a:(b, c) => b + c}")
+          );
+
+          it(
+            'should not be fooled by no-parentheses arrow functions',
+            testRoundTrip("{a:a => a + 1}")
+          );
+
+          it('should stringify extracted methods', function () {
+            var fn = eval('({ foo(x) { return x + 1; } })').foo;
+            expect(stringify(fn)).to.equal('function foo(x) { return x + 1; }');
+          });
+
+          it('should stringify extracted generators', function () {
+            var fn = eval('({ *foo(x) { yield x; } })').foo;
+            expect(stringify(fn)).to.equal('function* foo(x) { yield x; }');
+          });
+
+          // It's difficult to disambiguate between this and the arrow function case. Since the latter is probably
+          // much more common than this pattern (who creates empty-named methods ever?), we don't even try. But this
+          // test is here as documentation of a known limitation of this feature.
+          it.skip('should stringify extracted methods with empty names', function () {
+            var fn = eval('({ ""(x) { return x + 1; } })')[''];
+            expect(stringify(fn)).to.equal('function (x) { return x + 1; }');
+          });
+
+          it('should handle transplanted names', function () {
+            var fn = eval('({ foo(x) { return x + 1; } })').foo;
+            expect(stringify({ bar: fn })).to.equal('{bar:function foo(x) { return x + 1; }}');
+          });
+
+          it('should handle transplanted names with generators', function () {
+            var fn = eval('({ *foo(x) { yield x; } })').foo;
+            expect(stringify({ bar: fn })).to.equal('{bar:function* foo(x) { yield x; }}');
+          });
+
+          it(
+            'should reindent methods',
+            test(
+              eval(
+'               ({\n' +
+'                 fn() {\n' +
+'                   if (true) {\n' +
+'                     return "hello";\n' +
+'                   }\n' +
+'                 }\n' +
+'               })'),
+              '{\n  fn() {\n    if (true) {\n      return "hello";\n    }\n  }\n}',
+              2
+            )
+          );
+        });
       }
     });
 
