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)), "", "exec") + compiled_without = compile(parse(without_comments), "", "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():