diff --git a/README.md b/README.md index 0e99426..705c272 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ The only difference is there is a new type of tree node: Comment >>> tree.body[0].value '# comment to hello' >>> dump(tree) -"Module(body=[Comment(value='# comment to hello'), Assign(targets=[Name(id='hello', ctx=Store())], value=Constant(value='hello', kind=None), type_comment=None)], type_ignores=[])" +"Module(body=[Assign(targets=[Name(id='hello', ctx=Store())], value=Constant(value='hello', kind=None), type_comment=None), Comment(value='# comment to hello', inline=True)], type_ignores=[])" ``` If you have python3.9 or above it's also possible to unparse the tree object with its comments preserved. ``` @@ -35,21 +35,5 @@ hello = 'hello' ``` More examples can be found in test_parse.py and test_unparse.py. -## Notes -1. Right now it is assumed that there is no difference between inlined comments and regular. -All inlined comments become regular after the tree object is unparsed. - -2. Inlined comments for class- (def-, if-, ...) block shift "inside" body of the corresponding block: - ``` - >>> source = """class Foo: # c1 - ... pass - ... """ - >>> unparse(parse(source)) - >>> print(unparse(parse(source))) - class Foo: - # c1 - pass - ``` - ## Contributing You are welcome to open an issue or create a pull request \ No newline at end of file diff --git a/ast_comments.py b/ast_comments.py index d6ac38a..273b859 100644 --- a/ast_comments.py +++ b/ast_comments.py @@ -8,7 +8,11 @@ class Comment(ast.AST): value: str - _fields = ("value",) + inline: bool + _fields = ( + "value", + "inline", + ) _CONTAINER_ATTRS = ["body", "handlers", "orelse", "finalbody"] @@ -35,6 +39,7 @@ def _enrich(source: Union[str, bytes], tree: ast.AST) -> None: end_lineno, end_col_offset = t.end c = Comment( value=t.string, + inline=t.string != t.line.strip("\n").strip(" "), lineno=lineno, col_offset=col_offset, end_lineno=end_lineno, @@ -72,7 +77,7 @@ def _enrich(source: Union[str, bytes], tree: ast.AST) -> None: attr = getattr(target_node, target_attr) attr.append(c_node) - attr.sort(key=lambda x: (x.end_lineno, not isinstance(x, Comment))) + attr.sort(key=lambda x: (x.end_lineno, isinstance(x, Comment))) def _get_tree_intervals( @@ -97,6 +102,7 @@ def _get_tree_intervals( return res +# get min and max line from a source tree def _get_interval(items: List[ast.AST]) -> Tuple[int, int]: linenos, end_linenos = [], [] for item in items: @@ -110,7 +116,10 @@ def _get_interval(items: List[ast.AST]) -> Tuple[int, int]: class _Unparser(ast._Unparser): def visit_Comment(self, node: Comment) -> None: - self.fill(node.value) + if node.inline: + self.write(f" {node.value}") + else: + self.fill(node.value) def visit_If(self, node: ast.If) -> None: def _get_first_not_comment_idx(orelse: list[ast.stmt]) -> int: diff --git a/test_parse.py b/test_parse.py index 6ddd9fe..872a684 100644 --- a/test_parse.py +++ b/test_parse.py @@ -12,6 +12,7 @@ def test_single_comment_in_tree(): nodes = parse(source).body assert len(nodes) == 1 assert isinstance(nodes[0], Comment) + assert not nodes[0].inline def test_separate_line_single_line(): @@ -25,14 +26,16 @@ def test_separate_line_single_line(): nodes = parse(source).body assert len(nodes) == 2 assert isinstance(nodes[0], Comment) + assert not nodes[0].inline -def test_inline_comment_before_statement(): - """Inlined comment goes before statement.""" +def test_inline_comment_after_statement(): + """Inlined comment goes after statement.""" source = """hello = 'hello' # comment to hello""" nodes = parse(source).body assert len(nodes) == 2 - assert isinstance(nodes[0], Comment) + assert isinstance(nodes[1], Comment) + assert nodes[1].inline def test_separate_line_multiline(): @@ -50,6 +53,8 @@ def test_separate_line_multiline(): assert isinstance(nodes[1], Comment) assert nodes[0].value == "# comment to hello 1" assert nodes[1].value == "# comment to hello 2" + assert not nodes[0].inline + assert not nodes[1].inline def test_multiline_and_inline_combined(): @@ -64,8 +69,11 @@ def test_multiline_and_inline_combined(): ) nodes = parse(source).body assert nodes[0].value == "# comment to hello 1" + assert not nodes[0].inline assert nodes[1].value == "# comment to hello 2" - assert nodes[2].value == "# comment to hello 3" + assert not nodes[1].inline + assert nodes[3].value == "# comment to hello 3" + assert nodes[3].inline def test_unrelated_comment(): @@ -79,6 +87,7 @@ def test_unrelated_comment(): nodes = parse(source).body assert len(nodes) == 2 assert isinstance(nodes[1], Comment) + assert not nodes[1].inline def test_comment_to_function(): @@ -93,8 +102,10 @@ def foo(*args, **kwargs): nodes = parse(source).body assert len(nodes) == 2 assert nodes[0].value == "# comment to function 'foo'" + assert not nodes[0].inline function_node = nodes[1] - assert function_node.body[0].value == "# comment to print" + assert function_node.body[1].value == "# comment to print" + assert function_node.body[1].inline def test_comment_to_class(): @@ -114,10 +125,13 @@ def foo(self): assert len(nodes) == 2 assert nodes[0].value == "# comment to class 'Foo'" + assert not nodes[0].inline class_body = nodes[1].body - assert isinstance(class_body[0], Comment) - assert isinstance(class_body[1], ast.Assign) + assert isinstance(class_body[0], ast.Assign) + assert isinstance(class_body[1], Comment) + assert class_body[1].inline assert isinstance(class_body[2], Comment) + assert not class_body[2].inline assert isinstance(class_body[3], ast.FunctionDef) @@ -125,7 +139,8 @@ def test_parse_again(): """We can parse AstNode objects.""" source = """hello = 'hello' # comment to hello""" nodes = parse(parse(source)).body - assert isinstance(nodes[0], Comment) + assert isinstance(nodes[1], Comment) + assert nodes[1].inline def test_parse_ast(): @@ -142,13 +157,14 @@ def test_multiple_statements_in_line(): def test_comment_to_multiple_statements(): - """Comment goes before statements.""" + """Comment goes behind statements.""" source = """a=1; b=2 # hello""" nodes = parse(source).body assert len(nodes) == 3 - assert isinstance(nodes[0], Comment) + assert isinstance(nodes[0], ast.Assign) assert isinstance(nodes[1], ast.Assign) - assert isinstance(nodes[2], ast.Assign) + assert isinstance(nodes[2], Comment) + assert nodes[2].inline def test_comments_to_if(): @@ -168,6 +184,7 @@ def test_comments_to_if(): body_nodes = nodes[0].body assert len(body_nodes) == 2 assert isinstance(body_nodes[0], Comment) + assert body_nodes[0].inline assert isinstance(body_nodes[1], ast.Expr) orelse_nodes = nodes[0].orelse @@ -175,11 +192,13 @@ def test_comments_to_if(): orelse_if_nodes = orelse_nodes[0].body assert len(orelse_if_nodes) == 2 assert isinstance(orelse_if_nodes[0], Comment) + assert orelse_if_nodes[0].inline assert isinstance(orelse_if_nodes[1], ast.Expr) orelse_else_nodes = orelse_nodes[0].orelse assert len(orelse_else_nodes) == 2 assert isinstance(orelse_else_nodes[0], Comment) + assert orelse_else_nodes[0].inline assert isinstance(orelse_else_nodes[1], ast.Expr) @@ -198,14 +217,14 @@ def test_comments_to_for(): body_nodes = nodes[0].body assert len(body_nodes) == 3 assert isinstance(body_nodes[0], Comment) - assert isinstance(body_nodes[1], Comment) - assert isinstance(body_nodes[2], ast.Continue) + assert isinstance(body_nodes[1], ast.Continue) + assert isinstance(body_nodes[2], Comment) orelse_nodes = nodes[0].orelse assert len(orelse_nodes) == 3 assert isinstance(orelse_nodes[0], Comment) - assert isinstance(orelse_nodes[1], Comment) - assert isinstance(orelse_nodes[2], ast.Pass) + assert isinstance(orelse_nodes[1], ast.Pass) + assert isinstance(orelse_nodes[2], Comment) def test_comments_to_try(): @@ -230,8 +249,8 @@ def test_comments_to_try(): body_nodes = nodes[0].body assert len(body_nodes) == 3 assert isinstance(body_nodes[0], Comment) - assert isinstance(body_nodes[1], Comment) - assert isinstance(body_nodes[2], ast.Expr) + assert isinstance(body_nodes[1], ast.Expr) + assert isinstance(body_nodes[2], Comment) handlers_nodes = nodes[0].handlers assert len(handlers_nodes) == 2 @@ -239,20 +258,20 @@ def test_comments_to_try(): handler_nodes = h_node.body assert len(handler_nodes) == 3 assert isinstance(handler_nodes[0], Comment) - assert isinstance(handler_nodes[1], Comment) - assert isinstance(handler_nodes[2], ast.Pass) + assert isinstance(handler_nodes[1], ast.Pass) + assert isinstance(handler_nodes[2], Comment) else_nodes = nodes[0].orelse assert len(else_nodes) == 3 assert isinstance(else_nodes[0], Comment) - assert isinstance(else_nodes[1], Comment) - assert isinstance(else_nodes[2], ast.Expr) + assert isinstance(else_nodes[1], ast.Expr) + assert isinstance(else_nodes[2], Comment) finalbody_nodes = nodes[0].finalbody assert len(finalbody_nodes) == 3 assert isinstance(finalbody_nodes[0], Comment) - assert isinstance(finalbody_nodes[1], Comment) - assert isinstance(finalbody_nodes[2], ast.Expr) + assert isinstance(finalbody_nodes[1], ast.Expr) + assert isinstance(finalbody_nodes[2], Comment) def test_comment_to_multiline_expr(): @@ -267,5 +286,6 @@ def test_comment_to_multiline_expr(): if_node = parse(source).body[0] body_nodes = if_node.body assert len(body_nodes) == 2 - assert isinstance(body_nodes[0], Comment) - assert isinstance(body_nodes[1], ast.Expr) + assert isinstance(body_nodes[0], ast.Expr) + assert isinstance(body_nodes[1], Comment) + assert body_nodes[1].inline diff --git a/test_unparse.py b/test_unparse.py index 15005b5..f706e70 100644 --- a/test_unparse.py +++ b/test_unparse.py @@ -93,7 +93,7 @@ def test_comment_to_class(): """ # comment to class 'Foo' class Foo: - var = "Foo var" # comment to 'Foo.var' + var = "Foo var" # comment to 'Foo.var' # comment to method 'foo' def foo(self): @@ -134,10 +134,10 @@ def test_comments_to_for(): """Comments to for/else blocks.""" source = dedent( """ - for i in range(10): # for comment - continue # continue comment - else: # else comment - pass # pass comment + for i in range(10): # for comment + continue # continue comment + else: # else comment + pass # pass comment """ ) _test_unparse(source) @@ -147,16 +147,16 @@ def test_comments_to_try(): """Comments to try/except/else/finally blocks.""" source = dedent( """ - try: # try comment - 1 / 0 # expr comment + try: # try comment + 1 / 0 # expr comment except ValueError: # except ValueError comment - pass # pass comment - except KeyError: # except KeyError - pass # pass comment - else: # else comment - print() # print comment - finally: # finally comment - print() # print comment + pass # pass comment + except KeyError: # except KeyError + pass # pass comment + else: # else comment + print() # print comment + finally: # finally comment + print() # print comment """ ) _test_unparse(source) @@ -193,3 +193,23 @@ def test_nested_ifs_to_elifs(): assert r.count("if") == 2 # if and elif assert r.count("elif") == 1 assert r.count("else") == 1 + + +def test_inline_comment_stays_inline(): + source = dedent( + """ + # abc + + class Foo: + pass + + class Bar: # def + pass + + class FooBar: + # ghi + pass + """ + ) + _test_unparse(source) + assert source.strip("\n") == unparse(parse(source))