Make all comment options in cli available in js api

Also removing more code within "loop" while at it.

6 files changed, 103 insertions, 37 deletions

diff --git a/README.md b/README.md
index 4f5b21a0..e0aa4ba9 100644
--- a/README.md
+++ b/README.md
@@ -849,8 +849,10 @@ which we care about here are `source_map` and `comments`.
 #### Keeping comments in the output
 
 In order to keep certain comments in the output you need to pass the
-`comments` option.  Pass a RegExp or a function.  If you pass a RegExp, only
-those comments whose body matches the regexp will be kept.  Note that body
+`comments` option.  Pass a RegExp, boolean or a function.  Stringified options
+`all` and `some` can be passed too, where `some` behaves like it's cli
+equivalent `--comments` without passing a value. If you pass a RegExp,
+only those comments whose body matches the regexp will be kept.  Note that body
 means without the initial `//` or `/*`.  If you pass a function, it will be
 called for every comment in the tree and will receive two arguments: the
 node that the comment is attached to, and the comment token itself.
diff --git a/bin/uglifyjs b/bin/uglifyjs
index 3f0c8254..8d7bd759 100755
--- a/bin/uglifyjs
+++ b/bin/uglifyjs
@@ -250,25 +250,10 @@ if (ARGS.keep_fnames) {
 if (BEAUTIFY)
     UglifyJS.merge(OUTPUT_OPTIONS, BEAUTIFY);
 
-if (ARGS.comments != null) {
-    if (/^\/.*\/[a-zA-Z]*$/.test(ARGS.comments)) {
-        try {
-            OUTPUT_OPTIONS.comments = extractRegex(ARGS.comments);
-        } catch (e) {
-            print_error("ERROR: Invalid --comments: " + e.message);
-        }
-    } else if (ARGS.comments == "all") {
-        OUTPUT_OPTIONS.comments = true;
-    } else {
-        OUTPUT_OPTIONS.comments = function(node, comment) {
-            var text = comment.value;
-            var type = comment.type;
-            if (type == "comment2") {
-                // multiline comment
-                return /@preserve|@license|@cc_on/i.test(text);
-            }
-        }
-    }
+if (ARGS.comments === "") {
+    OUTPUT_OPTIONS.comments = "some";
+} else {
+    OUTPUT_OPTIONS.comments = ARGS.comments;
 }
 
 var files = ARGS._.slice();
diff --git a/lib/output.js b/lib/output.js
index 801f7516..a3c6b4ab 100644
--- a/lib/output.js
+++ b/lib/output.js
@@ -70,6 +70,49 @@ function OutputStream(options) {
         keep_quoted_props: false
     }, true);
 
+    // Convert comment option to RegExp if neccessary and set up comments filter
+    if (typeof options.comments === "string" && /^\/.*\/[a-zA-Z]*$/.test(options.comments)) {
+        var regex_pos = options.comments.lastIndexOf("/");
+        options.comments = new RegExp(
+            options.comments.substr(1, regex_pos - 1),
+            options.comments.substr(regex_pos + 1)
+        );
+    }
+    if (options.comments instanceof RegExp) {
+        options.comments = (function(f) {
+            return function(comment) {
+                return comment.type == "comment5" || f.test(comment.value);
+            }
+        })(options.comments);
+    }
+    else if (typeof options.comments === "function") {
+        options.comments = (function(f) {
+            return function(comment) {
+                return comment.type == "comment5" || f(this, comment);
+            }
+        })(options.comments);
+    }
+    else if (options.comments === "some") {
+        options.comments = function(comment) {
+            var text = comment.value;
+            var type = comment.type;
+            if (type == "comment2") {
+                // multiline comment
+                return /@preserve|@license|@cc_on/i.test(text);
+            }
+        }
+    }
+    else if (options.comments){ // NOTE includes "all" option
+        options.comments = function() {
+            return true;
+        }
+    } else {
+        // Falsy case, so reject all comments, except shebangs
+        options.comments = function(comment) {
+            return comment.type == "comment5";
+        }
+    }
+
     var indentation = 0;
     var current_col = 0;
     var current_line = 1;
@@ -435,7 +478,7 @@ function OutputStream(options) {
 
     AST_Node.DEFMETHOD("add_comments", function(output){
         if (output._readonly) return;
-        var c = output.option("comments"), self = this;
+        var self = this;
         var start = self.start;
         if (start && !start._comments_dumped) {
             start._comments_dumped = true;
@@ -458,19 +501,7 @@ function OutputStream(options) {
                 }));
             }
 
-            if (!c) {
-                comments = comments.filter(function(comment) {
-                    return comment.type == "comment5";
-                });
-            } else if (c.test) {
-                comments = comments.filter(function(comment){
-                    return comment.type == "comment5" || c.test(comment.value);
-                });
-            } else if (typeof c == "function") {
-                comments = comments.filter(function(comment){
-                    return comment.type == "comment5" || c(self, comment);
-                });
-            }
+            comments = comments.filter(output.option("comments"), self);
 
             // Keep single line comments after nlb, after nlb
             if (!output.option("beautify") && comments.length > 0 &&
diff --git a/test/input/comments/filter.js b/test/input/comments/filter.js
new file mode 100644
index 00000000..c752080f
--- /dev/null
+++ b/test/input/comments/filter.js
@@ -0,0 +1,3 @@
+// foo
+/*@preserve*/
+// bar
diff --git a/test/mocha/cli.js b/test/mocha/cli.js
index 38b57cd7..495b0076 100644
--- a/test/mocha/cli.js
+++ b/test/mocha/cli.js
@@ -2,11 +2,11 @@ var assert = require("assert");
 var exec = require("child_process").exec;
 
 describe("bin/uglifyjs", function () {
+    var uglifyjscmd = '"' + process.argv[0] + '" bin/uglifyjs';
     it("should produce a functional build when using --self", function (done) {
         this.timeout(5000);
 
-        var uglifyjs = '"' + process.argv[0] + '" bin/uglifyjs';
-        var command = uglifyjs + ' --self -cm --wrap WrappedUglifyJS';
+        var command = uglifyjscmd + ' --self -cm --wrap WrappedUglifyJS';
 
         exec(command, function (err, stdout) {
             if (err) throw err;
@@ -19,4 +19,34 @@ describe("bin/uglifyjs", function () {
             done();
         });
     });
+    it("Should be able to filter comments correctly with `--comment all`", function (done) {
+        var command = uglifyjscmd + ' test/input/comments/filter.js --comments all';
+
+        exec(command, function (err, stdout) {
+            if (err) throw err;
+
+            assert.strictEqual(stdout, "// foo\n/*@preserve*/\n// bar\n\n");
+            done();
+        });
+    });
+    it("Should be able to filter comments correctly with `--comment <RegExp>`", function (done) {
+        var command = uglifyjscmd + ' test/input/comments/filter.js --comments /r/';
+
+        exec(command, function (err, stdout) {
+            if (err) throw err;
+
+            assert.strictEqual(stdout, "/*@preserve*/\n// bar\n\n");
+            done();
+        });
+    });
+    it("Should be able to filter comments correctly with just `--comment`", function (done) {
+        var command = uglifyjscmd + ' test/input/comments/filter.js --comments';
+
+        exec(command, function (err, stdout) {
+            if (err) throw err;
+
+            assert.strictEqual(stdout, "/*@preserve*/\n\n");
+            done();
+        });
+    });
 });
diff --git a/test/mocha/comment-filter.js b/test/mocha/comment-filter.js
index ea2ec2eb..4330d1eb 100644
--- a/test/mocha/comment-filter.js
+++ b/test/mocha/comment-filter.js
@@ -7,6 +7,16 @@ describe("comment filters", function() {
         assert.strictEqual(ast.print_to_string({comments: /^!/}), "/*!test1*/\n//!test3\n//!test6\n//!test8\n");
     });
 
+    it("Should be able to filter comments with the 'all' option", function() {
+        var ast = UglifyJS.parse("/*!test1*/\n/*test2*/\n//!test3\n//test4\n<!--test5\n<!--!test6\n-->test7\n-->!test8");
+        assert.strictEqual(ast.print_to_string({comments: "all"}), "/*!test1*/\n/*test2*/\n//!test3\n//test4\n//test5\n//!test6\n//test7\n//!test8\n");
+    });
+
+    it("Should be able to filter commments with the 'some' option", function() {
+        var ast = UglifyJS.parse("// foo\n/*@preserve*/\n// bar\n/*@license*/\n//@license with the wrong comment type\n/*@cc_on something*/");
+        assert.strictEqual(ast.print_to_string({comments: "some"}), "/*@preserve*/\n/*@license*/\n/*@cc_on something*/\n");
+    });
+
     it("Should be able to filter comments by passing a function", function() {
         var ast = UglifyJS.parse("/*TEST 123*/\n//An other comment\n//8 chars.");
         var f = function(node, comment) {
@@ -16,6 +26,11 @@ describe("comment filters", function() {
         assert.strictEqual(ast.print_to_string({comments: f}), "/*TEST 123*/\n//8 chars.\n");
     });
 
+    it("Should be able to filter comments by passing regex in string format", function() {
+        var ast = UglifyJS.parse("/*!test1*/\n/*test2*/\n//!test3\n//test4\n<!--test5\n<!--!test6\n-->test7\n-->!test8");
+        assert.strictEqual(ast.print_to_string({comments: "/^!/"}), "/*!test1*/\n//!test3\n//!test6\n//!test8\n");
+    });
+
     it("Should be able to get the comment and comment type when using a function", function() {
         var ast = UglifyJS.parse("/*!test1*/\n/*test2*/\n//!test3\n//test4\n<!--test5\n<!--!test6\n-->test7\n-->!test8");
         var f = function(node, comment) {