support es6 arrow functions, fixes #389 (#393)

nodeca · Feb 15, 2018 · c62fd62 · c62fd62
1 parent bee7e99
commit c62fd62
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/lib/js-yaml/type/js/function.js b/lib/js-yaml/type/js/function.js
@@ -30,7 +30,8 @@ function resolveJavascriptFunction(data) {
     if (ast.type                    !== 'Program'             ||
         ast.body.length             !== 1                     ||
         ast.body[0].type            !== 'ExpressionStatement' ||
-        ast.body[0].expression.type !== 'FunctionExpression') {
+        (ast.body[0].expression.type !== 'ArrowFunctionExpression' &&
+          ast.body[0].expression.type !== 'FunctionExpression')) {
       return false;
     }
 
@@ -51,7 +52,8 @@ function constructJavascriptFunction(data) {
   if (ast.type                    !== 'Program'             ||
       ast.body.length             !== 1                     ||
       ast.body[0].type            !== 'ExpressionStatement' ||
-      ast.body[0].expression.type !== 'FunctionExpression') {
+      (ast.body[0].expression.type !== 'ArrowFunctionExpression' &&
+        ast.body[0].expression.type !== 'FunctionExpression')) {
     throw new Error('Failed to resolve function');
   }
 

diff --git a/test/samples-common/construct-javascript-function.js b/test/samples-common/construct-javascript-function.js
@@ -9,6 +9,8 @@ function testHandler(actual) {
 
   assert.strictEqual(actual.length, expected.length);
 
+  assert.strictEqual(actual.length, expected.length);
+
   assert.strictEqual(
     actual[0](),
     expected[0]());
@@ -26,6 +28,9 @@ testHandler.expected = [
   function () {
     return 42;
   },
+  function () {
+    return 72;
+  },
   function (x, y) {
     return x + y;
   },

diff --git a/test/samples-common/construct-javascript-function.yml b/test/samples-common/construct-javascript-function.yml
@@ -1,4 +1,5 @@
 - !!js/function 'function () { return 42 }'
+- !!js/function '() => { return 72 }'
 - !!js/function 'function (x, y) { return x + y; } '
 - !!js/function |
   function (foo) {
@@ -9,4 +10,4 @@
       second: 'sum',
       third: result
     };
-  }
+  }