From 6a98cfe33d44c688d32e4a333aa9540fd2142530 Mon Sep 17 00:00:00 2001
From: Bill Huneke <wahuneke@gmail.com>
Date: Thu, 26 Oct 2023 17:23:07 -0400
Subject: [PATCH] Add a utility function to strip Comment nodes from a parsed
 tree.

Resolves issue #23
---
 README.md       |  3 +++
 ast_comments.py | 13 +++++++++++++
 test_parse.py   | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 1b9622f..76fbe6c 100644
--- a/README.md
+++ b/README.md
@@ -32,6 +32,9 @@ If you have python3.9 or above it's also possible to unparse the tree object wit
 >>> print(unparse(tree))
 hello = 'hello'  # comment to hello
 ```
+**Note**: Python `compile()` cannot be run on the tree output from parse. The included `pre_compile_fixer()` function can be used to fix the tree (stripping 
+comment nodes) if it will be necessary to compile the output.
+
 More examples can be found in test_parse.py and test_unparse.py.
 
 ## Contributing
diff --git a/ast_comments.py b/ast_comments.py
index 809bfa4..eb005a9 100644
--- a/ast_comments.py
+++ b/ast_comments.py
@@ -262,3 +262,16 @@ def _get_first_not_comment_idx(orelse: list[ast.stmt]) -> int:
 
     def unparse(ast_obj: ast.AST) -> str:
         return _Unparser().visit(ast_obj)
+
+
+def pre_compile_fixer(tree: ast.AST) -> ast.AST:
+    """
+    The parse output from ast_comments cannot compile (see issue #23). This function can be
+    run to fix the output so that it can compile.  This transformer strips out Comment nodes.
+    """
+
+    class RewriteComments(ast.NodeTransformer):
+        def visit_Comment(self, node: ast.AST) -> ast.AST:
+            return None
+
+    return RewriteComments().visit(tree)
diff --git a/test_parse.py b/test_parse.py
index e85f8f3..b7ce1e6 100644
--- a/test_parse.py
+++ b/test_parse.py
@@ -1,11 +1,41 @@
 """Tests for `ast_comments.parse`."""
 
 import ast
+import dis
 from textwrap import dedent
 
 import pytest
 
-from ast_comments import Comment, parse
+from ast_comments import Comment, parse, pre_compile_fixer
+
+
+def test_compile_parse_output():
+    """
+    Output of `ast_comments.parse` can be compiled after applying the compile_fix helper function.
+    The tree resulting from this compilation should be the same as the tree resulting from
+    compiling the same code without comments.
+    """
+    with_comments = dedent(
+        """
+        func(1, 2)  # comment
+        # Another comment
+        x = 5 * "s"
+        """
+    )
+    without_comments = dedent(
+        """
+        func(1, 2)
+
+        x = 5 * "s"
+        """
+    )
+
+    compiled_with = compile(pre_compile_fixer(parse(with_comments)), "<string>", "exec")
+    compiled_without = compile(parse(without_comments), "<string>", "exec")
+
+    assert list(dis.get_instructions(compiled_with)) == list(
+        dis.get_instructions(compiled_without)
+    )
 
 
 def test_comment_at_start_of_inner_block_getting_correctly_parsed():