diff --git a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py index 00e6c5b9..38fe0e37 100644 --- a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py +++ b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py @@ -1,6 +1,7 @@ from __future__ import annotations import sys +import textwrap from enum import Enum, auto from typing import TypeVar, Optional, Union, cast, List @@ -55,13 +56,16 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> return super().visit(tree, p) def pre_visit(self, tree: T, p: P) -> Optional[T]: - if isinstance(tree, (JavaSourceFile, Label, ArrayDimension, ClassDeclaration, ExpressionStatement)): + if isinstance(tree, (JavaSourceFile, Label, ArrayDimension, ClassDeclaration)): self.cursor.put_message("indent_type", self.IndentType.ALIGN) elif isinstance(tree, Block): self.cursor.put_message("indent_type", self.IndentType.INDENT) elif isinstance(tree, (DictLiteral, CollectionLiteral, NewArray, ComprehensionExpression)): self.cursor.put_message("indent_type", self.IndentType.CONTINUATION_INDENT if self._other.use_continuation_indent.collections_and_comprehensions else self.IndentType.INDENT) + elif isinstance(tree, ExpressionStatement): + self.cursor.put_message("indent_type", self.IndentType.INDENT + if self._is_doc_comment(tree, self.cursor) else self.IndentType.ALIGN) elif isinstance(tree, Expression): self.cursor.put_message("indent_type", self.IndentType.INDENT) @@ -310,6 +314,25 @@ def visit_container(self, container: Optional[JContainer[J2]], return container return JContainer(before, js, container.markers) + @staticmethod + def _is_doc_comment(expression_statement: ExpressionStatement, cursor: Cursor) -> bool: + expr = expression_statement.expression + return isinstance(expr, Literal) and isinstance(expr.value_source, str) and ( + (expr.value_source.startswith('"""') and expr.value_source.endswith('"""')) or + (expr.value_source.startswith("'''") and expr.value_source.endswith("'''"))) and \ + cursor.first_enclosing(Block) is not None + + def visit_expression_statement(self, expression_statement: ExpressionStatement, p: P) -> J: + if self._is_doc_comment(expression_statement, self.cursor): + prefix_before = len(expression_statement.prefix.last_whitespace.split("\n")[-1]) + stm = cast(ExpressionStatement, super().visit_expression_statement(expression_statement, p)) + literal = cast(Literal, stm.expression) + shift = len(stm.prefix.last_whitespace.split("\n")[-1]) - prefix_before + return stm.with_expression( + literal.with_value_source(textwrap.indent(str(literal.value_source), shift * " ")[shift:])) + + return super().visit_expression_statement(expression_statement, p) + def _indent_to(self, space: Space, column: int, space_location: Optional[Union[PySpace.Location, Space.Location]]) -> Space: s = space whitespace = s.whitespace diff --git a/rewrite/tests/python/all/format/tabs_and_indents_visitor_test.py b/rewrite/tests/python/all/format/tabs_and_indents_visitor_test.py index 20aba063..a6c81a33 100644 --- a/rewrite/tests/python/all/format/tabs_and_indents_visitor_test.py +++ b/rewrite/tests/python/all/format/tabs_and_indents_visitor_test.py @@ -852,20 +852,92 @@ def my_function(a, b): return None -def test_docstring_alignment(): +def test_string_literal_assignment(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) rewrite_run( # language=python python( ''' - def my_function(): + a = """ + This is a string that + should not be modified. """ - This is a docstring that + ''' + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_string_literal_assignment_in_function(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + ''' + def my_function(): + a = """ + This is a string that should align with the function body. """ return None ''', ''' + def my_function(): + a = """ + This is a string that + should align with the function body. + """ + return None + ''' + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_string_literal_comment(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + ''' + 1+1 + """ + This is a comment that + should not be modified. + """ + ''' + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_int_literal(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + ''' + 1 + ''' + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_docstring_alignment(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + ''' + def my_function(): + """ + This is a docstring that + should align with the function body. + """ + return None + ''', + ''' def my_function(): """ This is a docstring that