From 99c9596550eb0390c54708ad09ceb4cd35949388 Mon Sep 17 00:00:00 2001 From: Tim Hatch Date: Sat, 3 Oct 2020 10:06:15 -0700 Subject: [PATCH] Ensure first lines remain intact with AddMissingHeaderRule --- fixit/rules/add_file_header.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/fixit/rules/add_file_header.py b/fixit/rules/add_file_header.py index 99b7ef72..0b6ddc3d 100644 --- a/fixit/rules/add_file_header.py +++ b/fixit/rules/add_file_header.py @@ -94,7 +94,18 @@ class AddMissingHeaderRule(CstLintRule): } ), expected_replacement="# header line\n# wrong header", - ) + ), + Invalid( + """ + #!/usr/bin/env python + # wrong header""", + config=LintConfig( + rule_config={ + "AddMissingHeaderRule": {"path": "*.py", "header": "# header line"} + } + ), + expected_replacement="#!/usr/bin/env python\n# header line\n# wrong header", + ), ] def __init__(self, context: CstContext) -> None: @@ -133,9 +144,14 @@ def visit_Module(self, node: cst.Module) -> None: if self.rule_disabled: return if not m.matches(node, m.Module(header=[*self.header_matcher, m.ZeroOrMore()])): + shebang, header = [], node.header + if header: + comment = header[0].comment + if isinstance(comment, cst.Comment) and comment.value.startswith("#!"): + shebang, header = [header[0]], header[1:] self.report( node, replacement=node.with_changes( - header=[*self.header_replacement, *node.header] + header=[*shebang, *self.header_replacement, *header] ), )