diff --git a/ng-annotate-main.js b/ng-annotate-main.js index b25c4e8..21e99f4 100644 --- a/ng-annotate-main.js +++ b/ng-annotate-main.js @@ -73,9 +73,30 @@ function matchDirectiveReturnObject(path) { // only matches inside directives // return { .. controller: function($scope, $timeout), ...} - return limit("directive", - (t.isReturnStatement(node) && node.argument && t.isObjectExpression(node.argument) && matchProp("controller", (path.get && path.get("argument.properties") || node.argument.properties))) || - (t.isArrowFunctionExpression(node) && node.body && t.isObjectExpression(node.body) && matchProp("controller", (path.get && path.get("body.properties") || node.body.properties)))); + var returnPath; + if (t.isReturnStatement(node) && node.argument) { + if (t.isObjectExpression(node.argument)) { + returnPath = matchProp("controller", (path.get && path.get("argument.properties") || node.argument.properties)); + } else if (path.get && t.isIdentifier(path.get("argument"))) { + var binding = path.scope.getBinding(node.argument.name); + var bound = binding && binding.path; + if (bound && t.isVariableDeclarator(bound)) { + var init = bound.get("init"); + if (init && t.isObjectExpression(init)) { + returnPath = matchProp("controller", init.get("properties")); + } + } + } + } + if (!returnPath) { + returnPath = + t.isArrowFunctionExpression(node) + && node.body + && t.isObjectExpression(node.body) + && matchProp("controller", (path.get && path.get("body.properties") || node.body.properties)); + } + + return limit("directive", returnPath); } function limit(name, path) { diff --git a/tests/references.js b/tests/references.js index d5e882b..13d55be 100644 --- a/tests/references.js +++ b/tests/references.js @@ -209,6 +209,39 @@ module.exports = { angular.module("test.feedback.pkg", []) .component("testFeedback", testFeedback); } + }, + { + name: "directive definition as a variable", + implicit: true, + input: function () { + function testDirective() { + var directiveDefinition = { + controller: testFeedbackController + }; + + return directiveDefinition; + + function testFeedbackController(foo) { + } + } + + angular.module("MyMod").directive("testDirective", testDirective); + }, + expected: function () { + function testDirective() { + testFeedbackController.$inject = ["foo"]; + var directiveDefinition = { + controller: testFeedbackController + }; + + return directiveDefinition; + + function testFeedbackController(foo) { + } + } + + angular.module("MyMod").directive("testDirective", testDirective); + } } ] }