diff --git a/src/tests/fixers_test.py b/src/tests/fixers_test.py index 6ad53dc..a5a6695 100644 --- a/src/tests/fixers_test.py +++ b/src/tests/fixers_test.py @@ -49,10 +49,17 @@ def setup(self): self.fixer = RaiseWithoutCauseFixer() def create_raise_no_cause_violation( - self, line: int, except_line: int = -1, exception_name: Optional[str] = None + self, + line: int, + except_line: int = -1, + exception_name: Optional[str] = None, + end_lineno: Optional[int] = None, ): code, _ = codes.RERAISE_NO_CAUSE - node_mock = MagicMock(spec=ast.Raise, lineno=line) + if not end_lineno: + end_lineno = line + + node_mock = MagicMock(spec=ast.Raise, lineno=line, end_lineno=end_lineno) except_node_mock = MagicMock(spec=ast.ExceptHandler, lineno=except_line) return RaiseWithoutCauseViolation( code, line, 0, "msg", "filename", node_mock, except_node_mock, exception_name @@ -83,3 +90,19 @@ def test_with_exception_name(self): assert_ast_is_valid(results) assert_unmodified_lines(lines, results, offending_offset) assert results[offending_offset].endswith("raise MyException() from error\n") + + def test_multiline_raise(self): + lines = read_sample_lines("except_reraise_no_cause", dir="autofix") + expected_modified_offsets = (38, 45) + offending_offset = 40 + dependent_offset, ending_modified_offset = expected_modified_offsets + violation = self.create_raise_no_cause_violation( + offending_offset + 1, dependent_offset + 1, end_lineno=ending_modified_offset + 1 + ) + + results = self.fixer.perform_fix(lines, violation) + + assert_ast_is_valid(results) + assert_unmodified_lines(lines, results, *expected_modified_offsets) + assert results[dependent_offset].endswith("except Exception as ex:\n") + assert results[ending_modified_offset].endswith(") from ex # multiline end!\n") diff --git a/src/tests/samples/autofix/except_reraise_no_cause.py b/src/tests/samples/autofix/except_reraise_no_cause.py index 5210fb6..ba09145 100644 --- a/src/tests/samples/autofix/except_reraise_no_cause.py +++ b/src/tests/samples/autofix/except_reraise_no_cause.py @@ -43,4 +43,4 @@ def longer_reraise(): 2, 3, 4, - ) + ) # multiline end! diff --git a/src/tryceratops/fixers/exception_block.py b/src/tryceratops/fixers/exception_block.py index a8151d8..c773e67 100644 --- a/src/tryceratops/fixers/exception_block.py +++ b/src/tryceratops/fixers/exception_block.py @@ -99,10 +99,19 @@ def _fix_except_handler(self, all_lines: List[str], offending_node: ast.ExceptHa def _fix_raise_no_cause( self, all_lines: List[str], violation: RaiseWithoutCauseViolation, exception_name: str ): - # TODO: What if this raise is multi-line? - guilty_line = all_lines[violation.line - 1] - new_line = re.sub(r"raise (.*)", rf"raise \1 from {exception_name}", guilty_line) - all_lines[violation.line - 1] = new_line + endline = violation.node.end_lineno or violation.line + is_singleline = violation.line == endline + + if is_singleline: + fix_offset = violation.line - 1 + guilty_line = all_lines[fix_offset] + new_line = re.sub(r"raise (.*)", rf"raise \1 from {exception_name}", guilty_line) + else: + fix_offset = endline - 1 + guilty_line = all_lines[fix_offset] + new_line = re.sub(r"(\))(.*)", rf"\1 from {exception_name}\2", guilty_line) + + all_lines[fix_offset] = new_line def perform_fix(self, lines: List[str], violation: RaiseWithoutCauseViolation) -> List[str]: all_lines = lines[:]