Support ES6 classes

- Annotated class
- Annotated expression class
- Annotated constructor
- Prologue directive on constructor

Originally authored by @nevcos (Filipe Costa <fcosta@fuze.com>) - closes #5.

Updated and improved by @FRSgit (kkkubas@o2.pl <pozdro1994elo>) - closes #13.

Fixes #4.

README.md

@@ -20,6 +20,10 @@ This fork contains the following changes:
 - Added support for ngInject in `export [default] function functionName() {...}`
   and `export [default] var varName = function [functionName]() {...}`.
 
+- Added support for ES6 classes with explicit `ngInject` annotations.
+  The support may not be perfect yet. For more information please see
+  [ES6 test file](tests/es6-classes.js).
+
 - Added support for dynamic `import()` syntax. If you use Webpack or a similar
   module loader you would probably like to compile to `esnext` modules for
   dynamic import support. To do that you will need to pass the

src/ng-annotate-main.js

@@ -469,25 +469,25 @@ function replaceNodeWith(node, newNode) {
     assert(done);
 }
 
-function insertArray(ctx, functionExpression, fragments, quot) {
+function insertArray(ctx, functionExpression, positioningNode, fragments, quot) {
     const args = stringify(ctx, functionExpression.params, quot);
 
     fragments.push({
-        start: functionExpression.range[0],
-        end: functionExpression.range[0],
+        start: positioningNode.range[0],
+        end: positioningNode.range[0],
         str: args.slice(0, -1) + ", ",
         loc: {
-            start: functionExpression.loc.start,
-            end: functionExpression.loc.start
+            start: positioningNode.loc.start,
+            end: positioningNode.loc.start
         }
     });
     fragments.push({
-        start: functionExpression.range[1],
-        end: functionExpression.range[1],
+        start: positioningNode.range[1],
+        end: positioningNode.range[1],
         str: "]",
         loc: {
-            start: functionExpression.loc.end,
-            end: functionExpression.loc.end
+            start: positioningNode.loc.end,
+            end: positioningNode.loc.end
         }
     });
 }
@@ -568,8 +568,19 @@ function judgeSuspects(ctx) {
         }
     }
 
+    const interimSuspects = suspects.map(function(node) {
+        if (isConstructorWithArgs(node)) {
+            while (node = node.$parent) {
+                if (node.type === "ExpressionStatement" || isClassExpression(node) || isClassDeclaration(node)) break;
+            }
+
+            node.$chained = chainedRegular;
+        }
+        return node;
+    });
+
     // create final suspects by jumping, following, uniq'ing, blocking
-    const finalSuspects = makeUnique(suspects.map(function(target) {
+    const finalSuspects = makeUnique(interimSuspects.map(function(target) {
         const jumped = jumpOverIife(target);
         const jumpedAndFollowed = followReference(jumped) || jumped;
 
@@ -589,12 +600,16 @@ function judgeSuspects(ctx) {
             return;
         }
 
+        let constructor;
+
         if (mode === "rebuild" && isAnnotatedArray(target)) {
             replaceArray(ctx, target, fragments, quot);
         } else if (mode === "remove" && isAnnotatedArray(target)) {
             removeArray(target, fragments);
         } else if (["add", "rebuild"].includes(mode) && isFunctionExpressionWithArgs(target)) {
-            insertArray(ctx, target, fragments, quot);
+            insertArray(ctx, target, target, fragments, quot);
+        } else if (["add", "rebuild"].includes(mode) && isClassExpression(target) && (constructor = findClassConstructorWithArgs(target))) {
+            insertArray(ctx, constructor.value, target, fragments, quot);
         } else if (isGenericProviderName(target)) {
             renameProviderDeclarationSite(ctx, target, fragments);
         } else {
@@ -800,7 +815,35 @@ function judgeInjectArraySuspect(node, ctx) {
 
     node = jumpOverIife(node);
 
-    if (ctx.isFunctionExpressionWithArgs(node)) {
+    let constructor;
+
+    if ((isClassExpression(node) || isClassDeclaration(node)) && (constructor = ctx.findClassConstructorWithArgs(node))) {
+        // /*@ngInject*/ class Foo { constructor($scope) {} }
+        // /*@ngInject*/ Foo = class { constructor($scope) {} }
+
+        const className = node.id ? node.id.name : declaratorName;
+        assert(className);
+
+        addRemoveInjectArray(
+            constructor.value.params,
+            insertPos,
+            className);
+
+    } else if (node.type === "ExpressionStatement" && node.expression.type === "AssignmentExpression" &&
+               isClassExpression(node.expression.right) && (constructor = ctx.findClassConstructorWithArgs(node.expression.right))) {
+        // foo.bar[0] = /*@ngInject*/ class($scope) {}
+
+        const className = ctx.srcForRange(node.expression.left.range);
+
+        addRemoveInjectArray(
+            constructor.value.params,
+            isSemicolonTerminated ? insertPos : {
+                pos: node.expression.right.range[1],
+                loc: node.expression.right.loc.end
+            },
+            className);
+
+    } else if (ctx.isFunctionExpressionWithArgs(node)) {
         // var x = 1, y = function(a,b) {}, z;
 
         assert(declaratorName);
@@ -1007,6 +1050,16 @@ function isAnnotatedArray(node) {
 
     return true;
 }
+
+function isConstructorWithArgs (node) {
+    return node.kind === 'constructor' && node.value.params.length >= 1;
+}
+function isClassExpression(node) {
+  return node.type === "ClassExpression";
+}
+function isClassDeclaration(node) {
+  return node.type === "ClassDeclaration";
+}
 function isFunctionExpressionWithArgs(node) {
     return node.type === "FunctionExpression" && node.params.length >= 1;
 }
@@ -1018,6 +1071,9 @@ function isFunctionDeclarationWithArgs(node) {
 function isGenericProviderName(node) {
     return node.type === "Literal" && typeof node.value === "string";
 }
+function findClassConstructorWithArgs(classFunction) {
+    return classFunction.body.body.find(isConstructorWithArgs);
+}
 
 function uniqifyFragments(fragments) {
     // must do in-place modification of ctx.fragments because shared reference
@@ -1164,9 +1220,12 @@ module.exports = function ngAnnotate(src, options) {
         suspects: suspects,
         blocked: blocked,
         lut: lut,
+        isClassExpression: isClassExpression,
+        isClassDeclaration: isClassDeclaration,
         isFunctionExpressionWithArgs: isFunctionExpressionWithArgs,
         isFunctionDeclarationWithArgs: isFunctionDeclarationWithArgs,
         isAnnotatedArray: isAnnotatedArray,
+        findClassConstructorWithArgs: findClassConstructorWithArgs,
         addModuleContextDependentSuspect: addModuleContextDependentSuspect,
         addModuleContextIndependentSuspect: addModuleContextIndependentSuspect,
         stringify: stringify,

src/nginject.js

@@ -12,8 +12,12 @@ module.exports = {
 function inspectNode(node, ctx) {
     if (node.type === "CallExpression") {
         inspectCallExpression(node, ctx);
+    } else if (node.$parent && node.$parent.type === "MethodDefinition") {
+        // Ignore method function, constructor is processed below
     } else if (node.type === "FunctionExpression" || node.type === "FunctionDeclaration") {
         inspectFunction(node, ctx);
+    } else if (node.type === "MethodDefinition" && node.kind === 'constructor') {
+        inspectConstructor(node, ctx);
     }
 }
 
@@ -76,6 +80,17 @@ function inspectFunction(node, ctx) {
     }
 }
 
+function inspectConstructor(node, ctx) {
+    const constructorFn = node.value;
+    const str = matchPrologueDirectives(ngAnnotatePrologueDirectives, constructorFn);
+    if (!str) {
+        return;
+    }
+    const block = (str === "ngNoInject");
+
+    addSuspect(node, ctx, block);
+}
+
 function matchPrologueDirectives(prologueDirectives, node) {
     const body = node.body.body;
 

tests/es6-classes.annotated.js

@@ -0,0 +1,84 @@
+// issue #3 (ng-annotate-patched) - Support for ES6 Classes
+
+(function(){
+    class ClassTest1 {
+        constructor($log) {}
+    }
+    /** @ngInject */
+    class ClassTest1_noargs {
+        constructor() {}
+    }
+    /** @ngInject */
+    class ClassTest1_annotated {
+        constructor($log) {}
+    }
+    ClassTest1_annotated.$inject = ["$log"];
+    class ClassTest1_annotated_constructor {
+        /** @ngInject */
+        constructor($log) {}
+    }
+    ClassTest1_annotated_constructor.$inject = ["$log"];
+    class ClassTest1_prologue_directive {
+        constructor($log) {
+            "ngInject";
+        }
+    }
+    ClassTest1_prologue_directive.$inject = ["$log"];
+
+    let ClassTest2 = class {
+        constructor($log) {}
+    };
+    /** @ngInject */
+    let ClassTest2_noargs = class {
+        constructor() {}
+    };
+    /** @ngInject */
+    let ClassTest2_annotated = class {
+        constructor($log) {}
+    };
+    ClassTest2_annotated.$inject = ["$log"];
+    let ClassTest2_annotated_expression = /** @ngInject */ ["$log", class {
+        constructor($log) {}
+    }];
+    let ClassTest2_annotated_constructor = ["$log", class {
+        /** @ngInject */
+        constructor($log) {}
+    }];
+    let ClassTest2_prologue_directive = ["$log", class {
+        constructor($log) {
+            "ngInject";
+        }
+    }];
+
+    let ClassTest3,
+        ClassTest3_noargs,
+        ClassTest3_annotated,
+        ClassTest3_annotated_expression,
+        ClassTest3_annotated_constructor,
+        ClassTest3_prologue_directive;
+
+    ClassTest3 = class {
+        constructor($log) {}
+    };
+    /** @ngInject */
+    ClassTest3_noargs = class {
+        constructor() {}
+    };
+    /** @ngInject */
+    ClassTest3_annotated = class {
+        constructor($log) {}
+    };
+    ClassTest3_annotated.$inject = ["$log"];
+    ClassTest3_annotated_expression = /** @ngInject */ ["$log", class {
+        constructor($log) {}
+    }];
+    ClassTest3_annotated_constructor = ["$log", class {
+        /** @ngInject */
+        constructor($log) {}
+    }];
+    ClassTest3_prologue_directive = ["$log", class {
+        constructor($log) {
+            "ngInject";
+        }
+    }];
+})();

tests/es6-classes.js

@@ -0,0 +1,79 @@
+// issue #3 (ng-annotate-patched) - Support for ES6 Classes
+
+(function(){
+    class ClassTest1 {
+        constructor($log) {}
+    }
+    /** @ngInject */
+    class ClassTest1_noargs {
+        constructor() {}
+    }
+    /** @ngInject */
+    class ClassTest1_annotated {
+        constructor($log) {}
+    }
+    class ClassTest1_annotated_constructor {
+        /** @ngInject */
+        constructor($log) {}
+    }
+    class ClassTest1_prologue_directive {
+        constructor($log) {
+            "ngInject";
+        }
+    }
+
+    let ClassTest2 = class {
+        constructor($log) {}
+    };
+    /** @ngInject */
+    let ClassTest2_noargs = class {
+        constructor() {}
+    };
+    /** @ngInject */
+    let ClassTest2_annotated = class {
+        constructor($log) {}
+    };
+    let ClassTest2_annotated_expression = /** @ngInject */ class {
+        constructor($log) {}
+    };
+    let ClassTest2_annotated_constructor = class {
+        /** @ngInject */
+        constructor($log) {}
+    };
+    let ClassTest2_prologue_directive = class {
+        constructor($log) {
+            "ngInject";
+        }
+    };
+
+    let ClassTest3,
+        ClassTest3_noargs,
+        ClassTest3_annotated,
+        ClassTest3_annotated_expression,
+        ClassTest3_annotated_constructor,
+        ClassTest3_prologue_directive;
+
+    ClassTest3 = class {
+        constructor($log) {}
+    };
+    /** @ngInject */
+    ClassTest3_noargs = class {
+        constructor() {}
+    };
+    /** @ngInject */
+    ClassTest3_annotated = class {
+        constructor($log) {}
+    };
+    ClassTest3_annotated_expression = /** @ngInject */ class {
+        constructor($log) {}
+    };
+    ClassTest3_annotated_constructor = class {
+        /** @ngInject */
+        constructor($log) {}
+    };
+    ClassTest3_prologue_directive = class {
+        constructor($log) {
+            "ngInject";
+        }
+    };
+})();

tests/run-tests.js

@@ -187,6 +187,12 @@ function run(ngAnnotate) {
     console.log("testing optionals/angular-dashboard-framework.js (removing annotations)");
     test(adf, ngAnnotate(adfAnnotated, {remove: true, plugin: [ngAnnotateAdfPlugin]}).src, "optionals/angular-dashboard-framework.js");
 
+    // issue #3 (ng-annotate-patched) - Support for ES6 Classes
+    console.log("testing es6 classes");
+    const es6Classes = slurp("tests/es6-classes.js");
+    const es6ClassesAnnotated = ngAnnotate(es6Classes, {add: true}).src;
+    test(slurp("tests/es6-classes.annotated.js"), es6ClassesAnnotated, "tests/es6-classes.annotated.js");
+
     console.log("testing performance");
     const ng1 = String(fs.readFileSync("tests/angular.js"));
     const ng5 = ng1 + ng1 + ng1 + ng1 + ng1;