Apply whitespace rules for LStrip and Trim to comment blocks

Author: hrdixon

src/main/java/com/hubspot/jinjava/tree/TreeParser.java

@@ -170,14 +170,17 @@ private Node getLastSibling() {
 
   private Node text(TextToken textToken) {
     if (interpreter.getConfig().isLstripBlocks()) {
-      if (scanner.hasNext() && scanner.peek().getType() == symbols.getTag()) {
-        textToken =
-          new TextToken(
-            StringUtils.stripEnd(textToken.getImage(), "\t "),
-            textToken.getLineNumber(),
-            textToken.getStartPosition(),
-            symbols
-          );
+      if (scanner.hasNext()) {
+        final int nextTokenType = scanner.peek().getType();
+        if (nextTokenType == symbols.getTag() || nextTokenType == symbols.getNote()) {
+          textToken =
+            new TextToken(
+              StringUtils.stripEnd(textToken.getImage(), "\t "),
+              textToken.getLineNumber(),
+              textToken.getStartPosition(),
+              symbols
+            );
+        }
       }
     }
 

src/main/java/com/hubspot/jinjava/tree/parse/TokenScanner.java

@@ -253,13 +253,18 @@ private Token newToken(int kind) {
       lastStart - lastNewlinePos + 1
     );
 
-    if (t instanceof TagToken) {
-      if (config.isTrimBlocks() && currPost < length && is[currPost] == '\n') {
-        lastNewlinePos = currPost;
-        ++currPost;
-        ++tokenStart;
-      }
+    if (
+      (t instanceof TagToken || t instanceof NoteToken) &&
+      config.isTrimBlocks() &&
+      currPost < length &&
+      is[currPost] == '\n'
+    ) {
+      lastNewlinePos = currPost;
+      ++currPost;
+      ++tokenStart;
+    }
 
+    if (t instanceof TagToken) {
       TagToken tt = (TagToken) t;
       if ("raw".equals(tt.getTagName())) {
         inRaw = 1;

src/test/java/com/hubspot/jinjava/tree/TreeParserTest.java

@@ -147,6 +147,18 @@ public void trimAndLstripBlocks() {
       .isEqualTo("<div>\n" + "        yay\n" + "</div>\n");
   }
 
+  @Test
+  public void trimAndLstripCommentBlocks() {
+    interpreter =
+      new Jinjava(
+        JinjavaConfig.newBuilder().withLstripBlocks(true).withTrimBlocks(true).build()
+      )
+        .newInterpreter();
+
+    assertThat(interpreter.render(parse("parse/tokenizer/whitespace-comment-tags.jinja")))
+      .isEqualTo("<div>\n" + "        yay\n" + "        whoop\n" + "</div>\n");
+  }
+
   @Test
   public void itWarnsAgainstMissingStartTags() {
     String expression = "{% if true %} foo {% endif %} {% endif %}";

src/test/java/com/hubspot/jinjava/tree/parse/TokenWhitespaceTest.java

@@ -21,6 +21,16 @@ public void trimBlocksTrimsAfterTag() {
     assertThat(tokens.get(2).getImage()).isEqualTo("        yay\n    ");
   }
 
+  @Test
+  public void trimBlocksTrimsAfterCommentTag() {
+    List<Token> tokens = scanTokens(
+      "parse/tokenizer/whitespace-comment-tags.jinja",
+      trimBlocksConfig()
+    );
+    assertThat(tokens.get(2).getImage()).isEqualTo("        yay\n    ");
+    assertThat(tokens.get(4).getImage()).isEqualTo("        whoop\n</div>\n");
+  }
+
   private List<Token> scanTokens(String srcPath, JinjavaConfig config) {
     try {
       return Lists.newArrayList(

src/test/resources/parse/tokenizer/whitespace-comment-tags.jinja

@@ -0,0 +1,8 @@
+<div>
+    {# a comment #}
+        yay
+    {# another comment
+       This time, over multiple lines
+     #}
+        whoop
+</div>