Fix and extra tests for literals

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

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