Skip to content

Commit

Permalink
Ensure first lines remain intact with AddMissingHeaderRule (#143)
Browse files Browse the repository at this point in the history
  • Loading branch information
thatch authored Oct 3, 2020
1 parent 1f11c18 commit fdabdb5
Showing 1 changed file with 18 additions and 2 deletions.
20 changes: 18 additions & 2 deletions fixit/rules/add_file_header.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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]
),
)

0 comments on commit fdabdb5

Please sign in to comment.