From 71680203b1ed66f6e288c1b0a4b7e93dfc457b0f Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Mon, 16 Dec 2024 14:36:57 +0100 Subject: [PATCH 01/16] Add test cases for tabs and indents visitor --- rewrite/rewrite/python/format/auto_format.py | 3 + .../python/format/tabs_and_indents_visitor.py | 17 ++ .../format/tabs_and_indents_visitor_test.py | 161 ++++++++++++++++++ 3 files changed, 181 insertions(+) create mode 100644 rewrite/rewrite/python/format/tabs_and_indents_visitor.py create mode 100644 rewrite/tests/python/all/format/tabs_and_indents_visitor_test.py diff --git a/rewrite/rewrite/python/format/auto_format.py b/rewrite/rewrite/python/format/auto_format.py index 9c115958..0ba58d50 100644 --- a/rewrite/rewrite/python/format/auto_format.py +++ b/rewrite/rewrite/python/format/auto_format.py @@ -6,6 +6,7 @@ from .normalize_tabs_or_spaces import NormalizeTabsOrSpacesVisitor from .remove_trailing_whitespace_visitor import RemoveTrailingWhitespaceVisitor from .spaces_visitor import SpacesVisitor +from .tabs_and_indents_visitor import TabsAndIndentsVisitor from .. import TabsAndIndentsStyle, GeneralFormatStyle from ..style import BlankLinesStyle, SpacesStyle, IntelliJ from ..visitor import PythonVisitor @@ -37,4 +38,6 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> tree = NormalizeLineBreaksVisitor(cu.get_style(GeneralFormatStyle) or GeneralFormatStyle(False), self._stop_after).visit(tree, p, self._cursor.fork()) tree = RemoveTrailingWhitespaceVisitor(self._stop_after).visit(tree, self._cursor.fork()) + tree = TabsAndIndentsVisitor(cu.get_style(TabsAndIndentsStyle) or IntelliJ.tabs_and_indents(), + self._stop_after).visit(tree, p, self._cursor.fork()) return tree diff --git a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py new file mode 100644 index 00000000..6467d699 --- /dev/null +++ b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from typing import TypeVar + +from rewrite import Tree +from rewrite.java import J +from rewrite.python import PythonVisitor, TabsAndIndentsStyle + +J2 = TypeVar('J2', bound=J) + + +class TabsAndIndentsVisitor(PythonVisitor): + + def __init__(self, style: TabsAndIndentsStyle, stop_after: Tree = None): + self._stop_after = stop_after + self._style = style + self._stop = False 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 new file mode 100644 index 00000000..d0ea494f --- /dev/null +++ b/rewrite/tests/python/all/format/tabs_and_indents_visitor_test.py @@ -0,0 +1,161 @@ +from rewrite.python import IntelliJ +from rewrite.python.format import TabsAndIndentsVisitor +from rewrite.test import rewrite_run, python, RecipeSpec, from_visitor + + +def test_multiline_call_with_post_arg_only(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + # noinspection PyInconsistentIndentation + rewrite_run( + # language=python + python( + """ + def long_function_name(var_one, var_two, + var_three, + var_four): + print(var_one) + """, + """ + def long_function_name(var_one, var_two, + var_three, + var_four): + print(var_one) + """ + ), + spec=RecipeSpec() + .with_recipes( + from_visitor(TabsAndIndentsVisitor(style)) + ) + ) + + +def test_multiline_call_with_args(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + """ + def example_function(): + result = some_method(10, 'foo', + another_arg=42, + final_arg="bar") + return result + """, + """ + def example_function(): + result = some_method(10, 'foo', + another_arg=42, + final_arg="bar") + return result + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_multiline_list(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + """ + def create_list(): + my_list = [ + 1, + 2, + 3, + 4 + ] + return my_list + """, + """ + def create_list(): + my_list = [ + 1, + 2, + 3, + 4 + ] + return my_list + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_nested_dictionary(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + """ + config = { + "section": { + "key1": "value1", + "key2": [10, 20, + 30], + }, + "another_section": {"nested_key": "val"} + } + """, + """ + config = { + "section": { + "key1": "value1", + "key2": [10, 20, + 30], + }, + "another_section": {"nested_key": "val"} + } + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_list_comprehension(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + """ + def even_numbers(n): + return [ x for x in range(n) + if x % 2 == 0] + """, + """ + def even_numbers(n): + return [ + x for x in range(n) + if x % 2 == 0 + ] + """ + ), + 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 + should align with the function body. + """ + return None + ''' + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) From 1fd9945808afdec9a70cf1570e86393935531ea5 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Thu, 19 Dec 2024 13:07:38 +0100 Subject: [PATCH 02/16] Add extra test cases for autoformat --- .../format/tabs_and_indents_visitor_test.py | 225 +++++++++++++++++- 1 file changed, 219 insertions(+), 6 deletions(-) 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 d0ea494f..fa26c2df 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 @@ -3,7 +3,192 @@ from rewrite.test import rewrite_run, python, RecipeSpec, from_visitor -def test_multiline_call_with_post_arg_only(): +def test_if_statement(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def check_value(x): + if x > 0: + return "Positive" + else: + return "Non-positive" + """, + """ + def check_value(x): + if x > 0: + return "Positive" + else: + return "Non-positive" + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_for_statement(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def sum_list(lst): + total = 0 + for num in lst: + total += num + return total + """, + """ + def sum_list(lst): + total = 0 + for num in lst: + total += num + return total + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_while_statement(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def countdown(n): + while n > 0: + print(n) + n -= 1 + """, + """ + def countdown(n): + while n > 0: + print(n) + n -= 1 + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_class_statement(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + class MyClass: + def __init__(self, value): + self.value = value + def get_value(self): + return self.value + """, + """ + class MyClass: + def __init__(self, value): + self.value = value + def get_value(self): + return self.value + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_with_statement(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def read_file(file_path): + with open(file_path, 'r') as file: + content = file.read() + return content + """, + """ + def read_file(file_path): + with open(file_path, 'r') as file: + content = file.read() + return content + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_try_statement(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def divide(a, b): + try: + return a / b + except ZeroDivisionError: + return None + """, + """ + def divide(a, b): + try: + return a / b + except ZeroDivisionError: + return None + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_basic_indent_modification(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(2).with_indent_size(4) + rewrite_run( + # language=python + python( + ''' + def my_function(): + return None + ''', + ''' + def my_function(): + return None + ''' + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_multiline_list(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size( + 4).with_continuation_indent(8) + rewrite_run( + # language=python + python( + """\ + my_list = [ + 1, + 2, + 3, + 4 + ] + """, + """\ + my_list = [ + 1, + 2, + 3, + 4 + ] + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_multiline_call_with_positional_args(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) # noinspection PyInconsistentIndentation rewrite_run( @@ -16,9 +201,37 @@ def long_function_name(var_one, var_two, print(var_one) """, """ - def long_function_name(var_one, var_two, - var_three, - var_four): + def long_function_name_2(var_one, var_two, + var_three, + var_four): + print(var_one) + """ + ), + spec=RecipeSpec() + .with_recipes( + from_visitor(TabsAndIndentsVisitor(style)) + ) + ) + + +def test_multiline_call_with_positional_args_and_no_arg_first_line(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + # noinspection PyInconsistentIndentation + rewrite_run( + # language=python + python( + """ + def long_function_name( + var_one, + var_two, var_three, + var_four): + print(var_one) + """, + """ + def long_function_name( + var_one, + var_two, var_three, + var_four): print(var_one) """ ), @@ -53,8 +266,8 @@ def example_function(): ) -def test_multiline_list(): - style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) +def test_multiline_list_inside_function(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) rewrite_run( # language=python python( From 52ae4a57b9c5ad8fce1a0cc192ce95dec5afe8a5 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Thu, 19 Dec 2024 16:56:15 +0100 Subject: [PATCH 03/16] More tests --- .../format/tabs_and_indents_visitor_test.py | 119 ++++++++++++++---- 1 file changed, 95 insertions(+), 24 deletions(-) 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 fa26c2df..4aa15bc8 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 @@ -51,6 +51,44 @@ def sum_list(lst): ) +def test_for_statement_with_list_comprehension(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def even_numbers(lst): + return [x for x in lst if x % 2 == 0] + """, + """ + def even_numbers(lst): + return [x for x in lst if x % 2 == 0] + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_for_statement_with_list_comprehension_multiline(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def even_numbers(lst): + return [x for x + in lst if x % 2 == 0] + """, + """ + def even_numbers(lst): + return [x for x + in lst if x % 2 == 0] + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + def test_while_statement(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) rewrite_run( @@ -188,7 +226,7 @@ def test_multiline_list(): ) -def test_multiline_call_with_positional_args(): +def test_multiline_call_with_positional_args_no_align_multiline(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) # noinspection PyInconsistentIndentation rewrite_run( @@ -196,14 +234,14 @@ def test_multiline_call_with_positional_args(): python( """ def long_function_name(var_one, var_two, - var_three, - var_four): + var_three, + var_four): print(var_one) """, """ - def long_function_name_2(var_one, var_two, - var_three, - var_four): + def long_function_name(var_one, var_two, + var_three, + var_four): print(var_one) """ ), @@ -214,6 +252,11 @@ def long_function_name_2(var_one, var_two, ) +def long_function_name(var_one, var_two, + var_three, + var_four): + print(var_one) + def test_multiline_call_with_positional_args_and_no_arg_first_line(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) # noinspection PyInconsistentIndentation @@ -242,24 +285,20 @@ def long_function_name( ) -def test_multiline_call_with_args(): +def test_multiline_call_with_args_without_multiline_align(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) rewrite_run( # language=python python( """ - def example_function(): - result = some_method(10, 'foo', - another_arg=42, - final_arg="bar") - return result + result = long_function_name(10, 'foo', + another_arg=42, + final_arg="bar") """, """ - def example_function(): - result = some_method(10, 'foo', - another_arg=42, - final_arg="bar") - return result + result = long_function_name(10, 'foo', + another_arg=42, + final_arg="bar") """ ), spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) @@ -273,13 +312,13 @@ def test_multiline_list_inside_function(): python( """ def create_list(): - my_list = [ - 1, - 2, - 3, - 4 - ] - return my_list + my_list = [ + 1, + 2, + 3, + 4 + ] + return my_list """, """ def create_list(): @@ -296,6 +335,38 @@ def create_list(): ) +def create_list(): + my_list = [ + 1, + 2, + 3, + 4 + ] + return my_list + + +def test_basic_dictionary(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + config = { + "key1": "value1", + "key2": "value2" + } + """, + """ + config = { + "key1": "value1", + "key2": "value2" + } + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + def test_nested_dictionary(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) rewrite_run( From dacd5b10b36e2678de4160ac4f468196e30660ba Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Thu, 19 Dec 2024 17:02:47 +0100 Subject: [PATCH 04/16] Add helper methods --- rewrite/rewrite/java/support_types.py | 38 +++++++++++++++++++++++++-- 1 file changed, 36 insertions(+), 2 deletions(-) diff --git a/rewrite/rewrite/java/support_types.py b/rewrite/rewrite/java/support_types.py index accff282..16d81c2b 100644 --- a/rewrite/rewrite/java/support_types.py +++ b/rewrite/rewrite/java/support_types.py @@ -106,8 +106,8 @@ def with_comments(self, comments: List[Comment]) -> Space: _whitespace: Optional[str] @property - def whitespace(self) -> Optional[str]: - return self._whitespace + def whitespace(self) -> str: + return self._whitespace if self._whitespace is not None else "" def with_whitespace(self, whitespace: Optional[str]) -> Space: return self if whitespace is self._whitespace else replace(self, _whitespace=whitespace) @@ -127,6 +127,40 @@ def format_first_prefix(cls, trees: List[J2], prefix: Space) -> List[J2]: return formatted_trees return trees + @property + def indent(self) -> str: + """ + The indentation after the last newline of either the last comment's suffix + or the global whitespace if no comments exist. + """ + return self._get_whitespace_indent(self.last_whitespace) + + @property + def last_whitespace(self) -> str: + """ + The raw suffix from the last comment if it exists, otherwise the global + whitespace (or empty string if whitespace is None). + """ + if self._comments: + return self._comments[-1].suffix + return self._whitespace if self._whitespace is not None else "" + + @staticmethod + def _get_whitespace_indent(whitespace: Optional[str]) -> str: + """ + A helper method that extracts everything after the last newline character + in `whitespace`. If no newline is present, returns `whitespace` as-is. + If the last newline is at the end, returns an empty string. + """ + # TODO: CAn be refactored probably to single line whitespace[whitespace.rfind('\n') + 1:] + if whitespace is None: + return "" + last_newline = whitespace.rfind('\n') + if last_newline >= 0: + return whitespace[last_newline + 1:] + elif last_newline == len(whitespace) - 1: + return "" + return whitespace EMPTY: ClassVar[Space] SINGLE_SPACE: ClassVar[Space] From d5bbcb6f1497b4b70327c0bba6dbe33bdc7a6e91 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Thu, 19 Dec 2024 17:03:44 +0100 Subject: [PATCH 05/16] Add helper method to cursor --- rewrite/rewrite/visitor.py | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/rewrite/rewrite/visitor.py b/rewrite/rewrite/visitor.py index 9d69274e..85b6dcc3 100644 --- a/rewrite/rewrite/visitor.py +++ b/rewrite/rewrite/visitor.py @@ -2,7 +2,7 @@ from abc import ABC from dataclasses import dataclass -from typing import TypeVar, Optional, Dict, List, Any, cast, Type, ClassVar, Generic +from typing import TypeVar, Optional, Dict, List, Any, cast, Type, ClassVar, Generic, Generator from .execution import RecipeRunException from .markers import Marker, Markers @@ -32,13 +32,13 @@ def put_message(self, key: str, value: object) -> None: object.__setattr__(self, 'messages', {}) self.messages[key] = value # type: ignore - def parent_tree_cursor(self) -> Optional[Cursor]: + def parent_tree_cursor(self) -> Cursor: c = self.parent while c is not None: - if isinstance(c.value, Tree): + if isinstance(c.value, Tree) or c.value == Cursor.ROOT_VALUE: return c c = c.parent - return None + raise ValueError("Expected to find parent tree cursor for " + str(self)) def first_enclosing_or_throw(self, type: Type[P]) -> P: result = self.first_enclosing(type) @@ -57,6 +57,30 @@ def first_enclosing(self, type_: Type[P]) -> Optional[P]: def fork(self) -> Cursor: return Cursor(self.parent.fork(), self.value) if self.parent else self + def get_path(self) -> Generator[Any]: + c = self + while c is not None: + yield c.value + c = c.parent + + def get_path_as_cursors(self) -> Generator[Cursor]: + c = self + while c is not None: + yield c + c = c.parent + + def get_nearest_message(self, key: str) -> Optional[object]: + for c in self.get_path_as_cursors(): + if c.messages is not None and key in c.messages: + return c.messages[key] + return None + + @property + def parent_or_throw(self) -> Cursor: + if self.parent is None: + raise ValueError("Cursor is expected to have a parent:", self) + return self.parent + class TreeVisitor(ABC, Generic[T, P]): _visit_count: int = 0 From dc7a4043e19841b1866ae1c5cb374b257eb70cb3 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Fri, 20 Dec 2024 14:36:46 +0100 Subject: [PATCH 06/16] Add WIP version of visitor --- rewrite/rewrite/java/support_types.py | 9 +- .../python/format/tabs_and_indents_visitor.py | 331 +++++++++++++++++- .../format/tabs_and_indents_visitor_test.py | 30 +- 3 files changed, 358 insertions(+), 12 deletions(-) diff --git a/rewrite/rewrite/java/support_types.py b/rewrite/rewrite/java/support_types.py index 16d81c2b..97aaef7d 100644 --- a/rewrite/rewrite/java/support_types.py +++ b/rewrite/rewrite/java/support_types.py @@ -152,15 +152,10 @@ def _get_whitespace_indent(whitespace: Optional[str]) -> str: in `whitespace`. If no newline is present, returns `whitespace` as-is. If the last newline is at the end, returns an empty string. """ - # TODO: CAn be refactored probably to single line whitespace[whitespace.rfind('\n') + 1:] - if whitespace is None: + if not whitespace: return "" last_newline = whitespace.rfind('\n') - if last_newline >= 0: - return whitespace[last_newline + 1:] - elif last_newline == len(whitespace) - 1: - return "" - return whitespace + return whitespace if last_newline == -1 else whitespace[last_newline + 1:] EMPTY: ClassVar[Space] SINGLE_SPACE: ClassVar[Space] diff --git a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py index 6467d699..9b528428 100644 --- a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py +++ b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py @@ -1,10 +1,16 @@ from __future__ import annotations -from typing import TypeVar +import sys +from enum import Enum, auto +from typing import TypeVar, Optional, Union, cast, List -from rewrite import Tree -from rewrite.java import J -from rewrite.python import PythonVisitor, TabsAndIndentsStyle +from rewrite import Tree, Cursor, list_map +from rewrite.java import J, Space, JRightPadded, JLeftPadded, JContainer, JavaSourceFile, EnumValueSet, Case, WhileLoop, \ + Block, If, ForLoop, ForEachLoop, Package, Import, Label, DoWhileLoop, ArrayDimension, ClassDeclaration, Empty, \ + Binary, MethodInvocation, FieldAccess, Identifier, Lambda +from rewrite.python import PythonVisitor, TabsAndIndentsStyle, PySpace, PyContainer, PyRightPadded, DictLiteral, \ + CollectionLiteral +from rewrite.visitor import P, T J2 = TypeVar('J2', bound=J) @@ -15,3 +21,320 @@ def __init__(self, style: TabsAndIndentsStyle, stop_after: Tree = None): self._stop_after = stop_after self._style = style self._stop = False + + def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> Optional[T]: + if parent is not None: + self._cursor = parent + if tree is None: + return cast(Optional[T], self.default_value(None, p)) + + for c in parent.get_path_as_cursors() if parent is not None else []: + v = c.value + print("V", v) + space = None + if isinstance(v, J): + space = v.prefix + elif isinstance(v, JRightPadded): + space = v.after + elif isinstance(v, JLeftPadded): + space = v.before + elif isinstance(v, JContainer): + space = v.before + + if space is not None and '\n' in space.last_whitespace: + indent = self.find_indent(space) + print("Indent: ", indent) + if indent != 0: + c.put_message("last_indent", indent) + + # TODO CHeck if this is 100% percent corrent + for v in parent.get_path() if parent is not None else []: + if isinstance(v, J): + self.pre_visit(v, p) + break + return super().visit(tree, p) + + def pre_visit(self, tree: T, p: P) -> Optional[T]: + if isinstance(tree, ( + JavaSourceFile, Package, Import, Label, DoWhileLoop, ArrayDimension, ClassDeclaration)): + self.cursor.put_message("indent_type", self.IndentType.ALIGN) + elif isinstance(tree, + (Block, If, If.Else, ForLoop, ForEachLoop, WhileLoop, Case, EnumValueSet, DictLiteral, + CollectionLiteral)): + # NOTE: Added DictLiteral here + self.cursor.put_message("indent_type", self.IndentType.INDENT) + else: + self.cursor.put_message("indent_type", self.IndentType.CONTINUATION_INDENT) + + return tree + + def post_visit(self, tree: T, p: P) -> Optional[T]: + if self._stop_after and tree == self._stop_after: + self._stop = True + return tree + + def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Location, Space.Location]], + p: P) -> Space: + if space is None: + return space # pyright: ignore [reportReturnType] + + self._cursor.put_message("last_location", loc) + parent = self._cursor.parent + + indent = cast(int, self.cursor.get_nearest_message("last_indent")) or 0 + indent_type = self.cursor.parent.get_nearest_message("indent_type") or self.IndentType.ALIGN + + if not space.comments and '\n' not in space.last_whitespace or parent is None: + return space + + cursor_value = self._cursor.value + + # Block spaces are always aligned to their parent + # The second condition ensure init blocks are ignored. + align_block_prefix_to_parent = loc is Space.Location.BLOCK_PREFIX and '\n' in space.whitespace and \ + (isinstance(cursor_value, Block) and not isinstance( + self.cursor.parent_tree_cursor().value, Block)) + + align_block_to_parent = loc in ( + Space.Location.BLOCK_END, + Space.Location.NEW_ARRAY_INITIALIZER_SUFFIX, + Space.Location.CATCH_PREFIX, + Space.Location.TRY_FINALLY, + Space.Location.ELSE_PREFIX, + ) + + if (loc == Space.Location.EXTENDS and "\n" in space.whitespace) or \ + Space.Location.EXTENDS == self.cursor.parent.get_message("last_location", None): + indent_type = self.IndentType.CONTINUATION_INDENT + + if align_block_prefix_to_parent or align_block_to_parent: + indent_type = self.IndentType.ALIGN + + if indent_type == self.IndentType.INDENT: + indent += self._style.indent_size + elif indent_type == self.IndentType.CONTINUATION_INDENT: + indent += self._style.continuation_indent + + s: Space = self._indent_to(space, indent, loc) + if isinstance(cursor_value, J) and not isinstance(cursor_value, EnumValueSet): + self.cursor.put_message("last_indent", indent) + elif loc == Space.Location.METHOD_SELECT_SUFFIX: + raise NotImplementedError("Method select suffix not implemented") + + return s + + # NOTE: INCOMPLETE + def visit_right_padded(self, right: Optional[JRightPadded[T]], + loc: Union[PyRightPadded.Location, JRightPadded.Location], p: P) -> Optional[ + JRightPadded[T]]: + + if right is None: + return None + + self.cursor = Cursor(self._cursor, right) + + indent: int = cast(int, self.cursor.get_nearest_message("last_indent")) or 0 + + t: T = right.element + after = right.after + # TODO: Check if the visit_and_cast is really required here + + if isinstance(t, J): + elem = t + if '\n' in right.after.last_whitespace or '\n' in elem.prefix.last_whitespace: + if loc in (JRightPadded.Location.FOR_CONDITION, + JRightPadded.Location.FOR_UPDATE): + raise ValueError("This case should not be possible, should be safe for removal...") + elif loc in (JRightPadded.Location.METHOD_DECLARATION_PARAMETER, + JRightPadded.Location.RECORD_STATE_VECTOR): + if isinstance(elem, Empty): + # NOTE: DONE + elem = elem.with_prefix(self._indent_to(elem.prefix, indent, loc.after_location)) + after = right.after + else: + container: JContainer[J] = cast(JContainer[J], self.cursor.parent_or_throw.value) + elements: List[J] = container.elements + last_arg: J = elements[-1] + + # TODO: style.MethodDeclarationParameters doesn't exist for Python + # but should be self._style.method_declaration_parameters.align_when_multiple + elem = self.visit_and_cast(elem, J, p) + after = self._indent_to(right.after, + indent if t is last_arg else self._style.continuation_indent, + loc.after_location) + + elif loc == JRightPadded.Location.METHOD_INVOCATION_ARGUMENT: + # NOTE: DONE + elem, after = self._visit_method_invocation_argument_j_type(elem, right, indent, loc, p) + elif loc in (JRightPadded.Location.NEW_CLASS_ARGUMENTS, + JRightPadded.Location.ARRAY_INDEX, + JRightPadded.Location.PARENTHESES, + JRightPadded.Location.TYPE_PARAMETER): + # NOTE: DONE + elem = self.visit_and_cast(elem, J, p) + after = self._indent_to(right.after, indent, loc.after_location) + elif loc == JRightPadded.Location.ANNOTATION_ARGUMENT: + raise NotImplementedError("Annotation argument not implemented") + else: + # NOTE: DONE + elem = self.visit_and_cast(elem, J, p) + after = self.visit_space(right.after, loc.after_location, p) + else: + # NOTE: Done + if loc in (JRightPadded.Location.NEW_CLASS_ARGUMENTS, JRightPadded.Location.METHOD_INVOCATION_ARGUMENT): + any_other_arg_on_own_line = False + if "\n" not in elem.prefix.last_whitespace: + # NOTE: Done + args: JContainer[J] = cast(JContainer[J], self.cursor.parent.value) + for arg in args.padding.elements: + if arg == self.cursor.value: + continue + if "\n" in arg.element.prefix.last_whitespace: + any_other_arg_on_own_line = True + break + if not any_other_arg_on_own_line: + elem = self.visit_and_cast(elem, J, p) + after = self._indent_to(right.after, indent, loc.after_location) + + if not any_other_arg_on_own_line: + # NOTE: Done + if not isinstance(elem, Binary): + # TODO SHould be able to merge + if not isinstance(elem, MethodInvocation): + self.cursor.put_message("last_indent", indent + self._style.continuation_indent) + elif "\n" in elem.prefix.last_whitespace: + self.cursor.put_message("last_indent", indent + self._style.continuation_indent) + else: + method_invocation = elem + select = method_invocation.select + if isinstance(select, (FieldAccess, Identifier, MethodInvocation)): + self.cursor.put_message("last_indent", indent + self._style.continuation_indent) + + elem = self.visit_and_cast(elem, J, p) + after = self.visit_space(right.after, loc.after_location, p) + else: + # NOTE: DONE + elem = self.visit_and_cast(elem, J, p) + after = self.visit_space(right.after, loc.after_location, p) + + t = cast(T, elem) + else: + after = self.visit_space(right.after, loc.after_location, p) + + self.cursor = self.cursor.parent + return right.with_after(after).with_element(t) + + # NOTE: INCOMPLETE + def visit_container(self, container: Optional[JContainer[J2]], + loc: Union[PyContainer.Location, JContainer.Location], p: P) -> JContainer[J2]: + self._cursor = Cursor(self._cursor, container) + if container is None: + return container + + indent = cast(int, self.cursor.get_nearest_message("last_indent")) or 0 + if '\n' in container.before.last_whitespace: + if loc in (JContainer.Location.TYPE_PARAMETERS, + JContainer.Location.IMPLEMENTS, + JContainer.Location.THROWS, + JContainer.Location.NEW_CLASS_ARGUMENTS): + before = self._indent_to(container.before, indent + self._style.continuation_indent, + loc.before_location) + self.cursor.put_message("indent_type", self.IndentType.ALIGN) + self.cursor.put_message("last_indent", indent + self._style.continuation_indent) + else: + before = self.visit_space(container.before, loc.before_location, p) + js = list_map(lambda t: self.visit_right_padded(t, loc.element_location, p), container.padding.elements) + else: + if loc in (JContainer.Location.TYPE_PARAMETERS, + JContainer.Location.IMPLEMENTS, + JContainer.Location.THROWS, + JContainer.Location.NEW_CLASS_ARGUMENTS, + JContainer.Location.METHOD_INVOCATION_ARGUMENTS): + self.cursor.put_message("indent_type", self.IndentType.CONTINUATION_INDENT) + before = self.visit_space(container.before, loc.before_location, p) + else: + before = self.visit_space(container.before, loc.before_location, p) + js = list_map(lambda t: self.visit_right_padded(t, loc.element_location, p), container.padding.elements) + + self._cursor = self._cursor.parent + + if container.padding.elements is js and container.before is before: + return container + return JContainer(before, js, container.markers) + + # NOTE: INCOMPLETE, Comments not supported + def _indent_to(self, space: Space, column: int, space_location: Space.Location) -> Space: + s = space + whitespace = s.whitespace + + if space_location == Space.Location.COMPILATION_UNIT_PREFIX and whitespace: + s = s.with_whitespace("") + elif not s.comments and "\n" not in s.last_whitespace: + return s + + if not s.comments: + indent = self.find_indent(s) + if indent != column: + shift = column - indent + s = s.with_whitespace(self._indent(whitespace, shift)) + else: + raise NotImplementedError("Comments not supported") + + return s + + def _indent(self, whitespace: str, shift: int): + return self._shift(whitespace, shift) + + def _shift(self, text: str, shift: int) -> str: + tab_indent = self._style.tab_size + if not self._style.use_tab_character: + tab_indent = sys.maxsize + + if shift > 0: + text += '\t' * (shift // tab_indent) + text += ' ' * (shift % tab_indent) + else: + if self._style.use_tab_character: + len_text = len(text) + (shift // tab_indent) + else: + len_text = len(text) + shift + if len_text >= 0: + text = text[:len_text] + + return text + + def find_indent(self, space: Space) -> int: + return self._get_length_of_whitespace(space.indent) + + def _get_length_of_whitespace(self, whitespace: Optional[str]) -> int: + if whitespace is None: + return 0 + length = 0 + for c in whitespace: + length += self._style.tab_size if c == '\t' else 1 + if c in ('\n', '\r'): + length = 0 + return length + + def _visit_method_invocation_argument_j_type(self, elem: J, right, indent, loc, p) -> tuple[J, Space]: + if "\n" not in elem.prefix.last_whitespace and isinstance(elem, Lambda): + body = elem.body + if not isinstance(body, Binary): + if "\n" not in body.prefix.last_whitespace: + self.cursor.parent_or_throw.put_message("last_indent", indent + self._style.continuation_indent) + + elem = self.visit_and_cast(elem, J, p) + after = self._indent_to(right.after, indent, loc.after_location) + if after.comments or "\n" in after.last_whitespace: + parent = self.cursor.parent_tree_cursor() + grandparent = parent.parent_tree_cursor() + # propagate indentation up in the method chain hierarchy + if isinstance(grandparent.value, MethodInvocation) and grandparent.value.select == parent.value: + grandparent.put_message("last_indent", indent) + grandparent.put_message("chained_indent", indent) + return elem, after + + class IndentType(Enum): + ALIGN = auto() + INDENT = auto() + CONTINUATION_INDENT = auto() 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 4aa15bc8..4004341e 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 @@ -3,7 +3,7 @@ from rewrite.test import rewrite_run, python, RecipeSpec, from_visitor -def test_if_statement(): +def test_if_else_statement(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) rewrite_run( # language=python @@ -27,6 +27,34 @@ def check_value(x): ) +def test_if_elif_else_statement(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def check_value(x): + if x > 0: + return "Positive" + elif x < 0: + return "Negative" + else: + return "Null" + """, + """ + def check_value(x): + if x > 0: + return "Positive" + elif x < 0: + return "Negative" + else: + return "Null" + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + def test_for_statement(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) rewrite_run( From 5312cef1aabc2462174d9f73dc327851b039a645 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Mon, 30 Dec 2024 13:04:30 +0100 Subject: [PATCH 07/16] Add extended test cases --- rewrite/rewrite/python/format/auto_format.py | 16 +- .../format/tabs_and_indents_visitor_test.py | 194 ++++++++++++++++-- 2 files changed, 186 insertions(+), 24 deletions(-) diff --git a/rewrite/rewrite/python/format/auto_format.py b/rewrite/rewrite/python/format/auto_format.py index 0ba58d50..f5e7c2fd 100644 --- a/rewrite/rewrite/python/format/auto_format.py +++ b/rewrite/rewrite/python/format/auto_format.py @@ -7,7 +7,8 @@ from .remove_trailing_whitespace_visitor import RemoveTrailingWhitespaceVisitor from .spaces_visitor import SpacesVisitor from .tabs_and_indents_visitor import TabsAndIndentsVisitor -from .. import TabsAndIndentsStyle, GeneralFormatStyle +from .wrapping_and_brances_visitor import WrappingAndBracesVisitor +from .. import TabsAndIndentsStyle, GeneralFormatStyle, WrappingAndBracesStyle from ..style import BlankLinesStyle, SpacesStyle, IntelliJ from ..visitor import PythonVisitor from ... import Recipe, Tree, Cursor @@ -29,15 +30,24 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> cu = tree if isinstance(tree, JavaSourceFile) else self._cursor.first_enclosing_or_throw(JavaSourceFile) tree = NormalizeFormatVisitor(self._stop_after).visit(tree, p, self._cursor.fork()) + tree = BlankLinesVisitor(cu.get_style(BlankLinesStyle) or IntelliJ.blank_lines(), self._stop_after).visit(tree, p, self._cursor.fork()) + + tree = WrappingAndBracesVisitor(WrappingAndBracesStyle(), self._stop_after).visit(tree, p, self._cursor.fork()) + tree = SpacesVisitor(cu.get_style(SpacesStyle) or IntelliJ.spaces(), self._stop_after).visit(tree, p, self._cursor.fork()) + tree = NormalizeTabsOrSpacesVisitor( cu.get_style(TabsAndIndentsStyle) or IntelliJ.tabs_and_indents(), self._stop_after ).visit(tree, p, self._cursor.fork()) + + tree = TabsAndIndentsVisitor(cu.get_style(TabsAndIndentsStyle) or IntelliJ.tabs_and_indents(), + self._stop_after).visit(tree, p, self._cursor.fork()) + tree = NormalizeLineBreaksVisitor(cu.get_style(GeneralFormatStyle) or GeneralFormatStyle(False), self._stop_after).visit(tree, p, self._cursor.fork()) + tree = RemoveTrailingWhitespaceVisitor(self._stop_after).visit(tree, self._cursor.fork()) - tree = TabsAndIndentsVisitor(cu.get_style(TabsAndIndentsStyle) or IntelliJ.tabs_and_indents(), - self._stop_after).visit(tree, p, self._cursor.fork()) + return tree 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 4004341e..ebe7faba 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 @@ -3,6 +3,29 @@ from rewrite.test import rewrite_run, python, RecipeSpec, from_visitor +def test_multi_assignment(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def assign_values(): + a, b = 1, 2 + x, y, z = 3, 4, 5 + return a, b, x, y, z + """, + """ + def assign_values(): + a, b = 1, 2 + x, y, z = 3, 4, 5 + return a, b, x, y, z + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + # spec = RecipeSpec().with_recipes(from_visitor(AutoFormatVisitor())) + ) + + def test_if_else_statement(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) rewrite_run( @@ -27,6 +50,81 @@ def check_value(x): ) +def test_if_else_statement_no_else_with_extra_statements(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def check_value(x): + if x > 0: + a = 1 + x + return "Positive" + a = -1 + x + return "Non-positive" + """, + """ + def check_value(x): + if x > 0: + a = 1 + x + return "Positive" + a = -1 + x + return "Non-positive" + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_if_else_statement_no_else_multi_return_values(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """\ + def check_value(x): + if x > 0: + a = 1 + x + return a, "Positive" + return a + """, + """\ + def check_value(x): + if x > 0: + a = 1 + x + return a, "Positive" + return a + """ + ), + # spec=RecipeSpec().with_recipes(from_visitor(AutoFormatVisitor())) + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_if_else_statement_no_else_multi_return_values_as_tuple(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """\ + def check_value(x): + if x > 0: + a = 1 + x + return (a, "Positive") + return a + """, + """\ + def check_value(x): + if x > 0: + a = 1 + x + return (a, "Positive") + return a + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + def test_if_elif_else_statement(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) rewrite_run( @@ -185,7 +283,7 @@ def read_file(file_path): ) -def test_try_statement(): +def test_try_statement_basic(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) rewrite_run( # language=python @@ -193,16 +291,46 @@ def test_try_statement(): """ def divide(a, b): try: - return a / b + c = a / b except ZeroDivisionError: return None + return c """, """ def divide(a, b): try: - return a / b + c = a / b except ZeroDivisionError: return None + return c + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_try_statement_with_multi_return(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def divide(a, b): + try: + c = a / b + if c > 42: return a, c + except ZeroDivisionError: + return None + return a, b + """, + """ + def divide(a, b): + try: + c = a / b + if c > 42: return a, c + except ZeroDivisionError: + return None + return a, b """ ), spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) @@ -234,7 +362,7 @@ def test_multiline_list(): # language=python python( """\ - my_list = [ + my_list = [ #cool 1, 2, 3, @@ -242,12 +370,12 @@ def test_multiline_list(): ] """, """\ - my_list = [ + my_list = [ #cool 1, 2, 3, 4 - ] + ] """ ), spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) @@ -280,11 +408,6 @@ def long_function_name(var_one, var_two, ) -def long_function_name(var_one, var_two, - var_three, - var_four): - print(var_one) - def test_multiline_call_with_positional_args_and_no_arg_first_line(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) # noinspection PyInconsistentIndentation @@ -363,16 +486,6 @@ def create_list(): ) -def create_list(): - my_list = [ - 1, - 2, - 3, - 4 - ] - return my_list - - def test_basic_dictionary(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) rewrite_run( @@ -447,6 +560,45 @@ def even_numbers(n): ) +def test_comment_alignment(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + ''' + # Informative comment 1 + def my_function(a, b): + # Informative comment 2 + + # Informative comment 3 + if a > b: + # cool + a = b + 1 + # cool + + return None # Informative comment 4 + # Informative comment 5 + ''', + ''' + # Informative comment 1 + def my_function(a, b): + # Informative comment 2 + + # Informative comment 3 + if a > b: + # cool + a = b + 1 + # cool + + return None # Informative comment 4 + # Informative comment 5 + ''' + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + return None + + def test_docstring_alignment(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) rewrite_run( From 63186b40ad192a76043ee272fbe4c690adc4c6d5 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Mon, 30 Dec 2024 13:36:30 +0100 Subject: [PATCH 08/16] WIP comment support and further extension --- .../python/format/tabs_and_indents_visitor.py | 77 ++++++++++++++++--- 1 file changed, 67 insertions(+), 10 deletions(-) diff --git a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py index 9b528428..2d7f73a0 100644 --- a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py +++ b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py @@ -7,7 +7,7 @@ from rewrite import Tree, Cursor, list_map from rewrite.java import J, Space, JRightPadded, JLeftPadded, JContainer, JavaSourceFile, EnumValueSet, Case, WhileLoop, \ Block, If, ForLoop, ForEachLoop, Package, Import, Label, DoWhileLoop, ArrayDimension, ClassDeclaration, Empty, \ - Binary, MethodInvocation, FieldAccess, Identifier, Lambda + Binary, MethodInvocation, FieldAccess, Identifier, Lambda, TextComment, Comment from rewrite.python import PythonVisitor, TabsAndIndentsStyle, PySpace, PyContainer, PyRightPadded, DictLiteral, \ CollectionLiteral from rewrite.visitor import P, T @@ -30,7 +30,6 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> for c in parent.get_path_as_cursors() if parent is not None else []: v = c.value - print("V", v) space = None if isinstance(v, J): space = v.prefix @@ -43,7 +42,6 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> if space is not None and '\n' in space.last_whitespace: indent = self.find_indent(space) - print("Indent: ", indent) if indent != 0: c.put_message("last_indent", indent) @@ -61,7 +59,7 @@ def pre_visit(self, tree: T, p: P) -> Optional[T]: elif isinstance(tree, (Block, If, If.Else, ForLoop, ForEachLoop, WhileLoop, Case, EnumValueSet, DictLiteral, CollectionLiteral)): - # NOTE: Added DictLiteral here + # NOTE: Added CollectionLiteral, DictLiteral here self.cursor.put_message("indent_type", self.IndentType.INDENT) else: self.cursor.put_message("indent_type", self.IndentType.CONTINUATION_INDENT) @@ -91,12 +89,12 @@ def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Locati # Block spaces are always aligned to their parent # The second condition ensure init blocks are ignored. + # TODO: Secon condition might be removed since it's not relevant for Python align_block_prefix_to_parent = loc is Space.Location.BLOCK_PREFIX and '\n' in space.whitespace and \ (isinstance(cursor_value, Block) and not isinstance( self.cursor.parent_tree_cursor().value, Block)) align_block_to_parent = loc in ( - Space.Location.BLOCK_END, Space.Location.NEW_ARRAY_INITIALIZER_SUFFIX, Space.Location.CATCH_PREFIX, Space.Location.TRY_FINALLY, @@ -123,6 +121,7 @@ def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Locati return s + # NOTE: INCOMPLETE def visit_right_padded(self, right: Optional[JRightPadded[T]], loc: Union[PyRightPadded.Location, JRightPadded.Location], p: P) -> Optional[ @@ -199,10 +198,7 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], if not any_other_arg_on_own_line: # NOTE: Done if not isinstance(elem, Binary): - # TODO SHould be able to merge - if not isinstance(elem, MethodInvocation): - self.cursor.put_message("last_indent", indent + self._style.continuation_indent) - elif "\n" in elem.prefix.last_whitespace: + if not isinstance(elem, MethodInvocation) or "\n" in elem.prefix.last_whitespace: self.cursor.put_message("last_indent", indent + self._style.continuation_indent) else: method_invocation = elem @@ -278,13 +274,74 @@ def _indent_to(self, space: Space, column: int, space_location: Space.Location) shift = column - indent s = s.with_whitespace(self._indent(whitespace, shift)) else: - raise NotImplementedError("Comments not supported") + def whitespace_indent(text: str) -> str: + # TODO: Placeholder function, taken from java openrewrite.StringUtils + indent = [] + for c in text: + if c == '\n' or c == '\r': + return ''.join(indent) + elif c.isspace(): + indent.append(c) + else: + return ''.join(indent) + return ''.join(indent) + + # TODO: This is the java version, however the python version is probably different + has_file_leading_comment = space.comments and ( + (space_location == Space.Location.COMPILATION_UNIT_PREFIX) or ( + space_location == Space.Location.BLOCK_END) or + (space_location == Space.Location.CLASS_DECLARATION_PREFIX and space.comments[0].multiline) + ) + + final_column = column + self._style.indent_size if space_location == Space.Location.BLOCK_END else column + last_indent: str = space.whitespace[space.whitespace.rfind('\n') + 1:] + indent = self._get_length_of_whitespace(whitespace_indent(last_indent)) + + if indent != final_column: + if (has_file_leading_comment or ("\n" in whitespace)) and ( + # Do not shift single-line comments at column 0. + not (s.comments and isinstance(s.comments[0], TextComment) and + not s.comments[0].multiline and self._get_length_of_whitespace(s.whitespace) == 0)): + shift = final_column - indent + s = s.with_whitespace(whitespace[:whitespace.rfind('\n') + 1] + self._indent(last_indent, shift)) + + final_space = s + last_comment_pos = len(s.comments) - 1 + + def _process_comment(i: int, c: Comment) -> Comment: + if isinstance(c, TextComment) and not c.multiline: + # Do not shift single line comments at col 0. + if i != last_comment_pos and self._get_length_of_whitespace(c.suffix) == 0: + return c + + prior_suffix = space.whitespace if i == 0 else final_space.comments[i - 1].suffix + + if space_location == Space.Location.BLOCK_END and i != len(final_space.comments) - 1: + to_column = column + self._style.indent_size + else: + to_column = column + new_c = c + if "\n" in prior_suffix or has_file_leading_comment: + new_c = self._indent_comment(c, prior_suffix, to_column) + + if '\n' in new_c.suffix: + suffix_indent = self._get_length_of_whitespace(new_c.suffix) + shift = to_column - suffix_indent + new_c = new_c.with_suffix(self._indent(new_c.suffix, shift)) + + return new_c + + s = s.with_comments(list_map(lambda i, c: _process_comment(c, i), s.comments)) return s def _indent(self, whitespace: str, shift: int): return self._shift(whitespace, shift) + def _indent_comment(self, comment: Comment, prior_suffix: str, to_column: int) -> Comment: + # TODO: Handle multiline here. + return comment + def _shift(self, text: str, shift: int) -> str: tab_indent = self._style.tab_size if not self._style.use_tab_character: From 62997a5fa8b88f2c2fd8c750e9aa33fde2e612c8 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Mon, 30 Dec 2024 13:43:13 +0100 Subject: [PATCH 09/16] Add Tree Printer --- rewrite/rewrite/python/utils/__init__.py | 5 + .../python/utils/tree_visiting_printer.py | 124 ++++++++++++++++++ 2 files changed, 129 insertions(+) create mode 100644 rewrite/rewrite/python/utils/__init__.py create mode 100644 rewrite/rewrite/python/utils/tree_visiting_printer.py diff --git a/rewrite/rewrite/python/utils/__init__.py b/rewrite/rewrite/python/utils/__init__.py new file mode 100644 index 00000000..ba1a43ca --- /dev/null +++ b/rewrite/rewrite/python/utils/__init__.py @@ -0,0 +1,5 @@ +from .tree_visiting_printer import * + +__all__ = [ + "TreeVisitingPrinter" +] diff --git a/rewrite/rewrite/python/utils/tree_visiting_printer.py b/rewrite/rewrite/python/utils/tree_visiting_printer.py new file mode 100644 index 00000000..4b6a2d6f --- /dev/null +++ b/rewrite/rewrite/python/utils/tree_visiting_printer.py @@ -0,0 +1,124 @@ +from typing import Optional, Union + +from rewrite import Cursor +from rewrite import Tree +from rewrite.java import Space, Literal, Identifier, JRightPadded, JLeftPadded, Modifier +from rewrite.python import PythonVisitor, PySpace +from rewrite.visitor import T, P + + +class TreeVisitingPrinter(PythonVisitor): + INDENT = " " + ELEMENT_PREFIX = "\\---" + CONTINUE_PREFIX = "|---" + UNVISITED_PREFIX = "#" + BRANCH_CONTINUE_CHAR = '|' + BRANCH_END_CHAR = '\\' + CONTENT_MAX_LENGTH = 120 + + _last_cursor_stack = [] + _lines = [] + + def __init__(self, indent: str = " "): + super().__init__() + self.INDENT = indent + + def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> Optional[T]: + if tree is None: + return super().visit(None, p, parent) # pyright: ignore [reportReturnType] + + _current_stack = list(self._cursor.get_path()) + _current_stack.reverse() + depth = len(_current_stack) + if not self._last_cursor_stack: + self._last_cursor_stack = _current_stack + [tree] + else: + diff_position = self.find_diff_pos(_current_stack, self._last_cursor_stack) + if diff_position >= 0: + for i in _current_stack[diff_position:]: + self._lines += [[depth, i]] + self._last_cursor_stack = self._last_cursor_stack[:diff_position] + + self._lines += [[depth, tree]] + self._last_cursor_stack = _current_stack + [tree] + return super().visit(tree, p, parent) # pyright: ignore [reportReturnType] + + def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Location, Space.Location]], + p: P) -> Space: + print("Loc", loc, "el,", self._print_element(self.cursor.parent.value)) + return super().visit_space(space, loc, p) + + def _print_tree(self) -> str: + output = "" + offset = 0 + for idx, (depth, element) in enumerate(self._lines): + offset = depth if idx == 0 else offset + padding = self.INDENT * (depth - offset) + if idx + 1 < len(self._lines) and self._lines[idx + 1][0] <= depth or idx + 1 == len(self._lines): + output += padding + self.CONTINUE_PREFIX + self._print_element(element) + "\n" + else: + output += padding + self.ELEMENT_PREFIX + self._print_element(element) + "\n" + return output + + def _print_element(self, element) -> str: + type_name = type(element).__name__ + line = [] + + if hasattr(element, "before"): + line.append(f"before= {self._print_space(element.before)}") + + if hasattr(element, "after"): + line.append(f"after= {self._print_space(element.after)}") + + if hasattr(element, "suffix"): + line.append(f"suffix= {self._print_space(element.suffix)}") + + if hasattr(element, "prefix"): + line.append(f"prefix= {self._print_space(element.prefix)}") + + if isinstance(element, Identifier): + type_name = f'{type_name} | "{element.simple_name}"' + + if isinstance(element, Literal): + type_name = f'{type_name} | {element.value_source}' + + if isinstance(element, JRightPadded): + return f'{type_name} | after= {self._print_space(element.after)}' + + if isinstance(element, JLeftPadded): + return f'{type_name} | before= {self._print_space(element.before)}' + + if isinstance(element, Modifier): + return type_name + ( + (" | " + element.type.name) if hasattr(element, "type") else "") + + if line: + return type_name + " | " + " | ".join(line) + return type_name + + @staticmethod + def _print_space(space: Space) -> str: + parts = [] + if space.whitespace: + parts.append(f'whitespace="{repr(space.whitespace)}"') + if space.comments: + parts.append(f'comments="{space.comments}"') + return " ".join(parts).replace("\n", "\\s\n") + + @staticmethod + def print_tree_all(tree: "Tree") -> str: + visitor = TreeVisitingPrinter() + visitor.visit(tree, None, None) + print(visitor._print_tree()) + return "" + + def find_diff_pos(self, cursor_stack, last_cursor_stack): + diff_pos = -1 + for i in range(len(cursor_stack)): + if i >= len(last_cursor_stack): + diff_pos = i + break + if cursor_stack[i] != last_cursor_stack[i]: + diff_pos = i + break + return diff_pos From 30c9550c30344f8bbcf907efc86607ed1b7c6358 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Wed, 8 Jan 2025 16:17:46 +0100 Subject: [PATCH 10/16] Fixes --- rewrite/rewrite/python/format/auto_format.py | 3 --- rewrite/rewrite/python/format/tabs_and_indents_visitor.py | 5 +++-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/rewrite/rewrite/python/format/auto_format.py b/rewrite/rewrite/python/format/auto_format.py index f5e7c2fd..54543512 100644 --- a/rewrite/rewrite/python/format/auto_format.py +++ b/rewrite/rewrite/python/format/auto_format.py @@ -7,7 +7,6 @@ from .remove_trailing_whitespace_visitor import RemoveTrailingWhitespaceVisitor from .spaces_visitor import SpacesVisitor from .tabs_and_indents_visitor import TabsAndIndentsVisitor -from .wrapping_and_brances_visitor import WrappingAndBracesVisitor from .. import TabsAndIndentsStyle, GeneralFormatStyle, WrappingAndBracesStyle from ..style import BlankLinesStyle, SpacesStyle, IntelliJ from ..visitor import PythonVisitor @@ -33,8 +32,6 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> tree = BlankLinesVisitor(cu.get_style(BlankLinesStyle) or IntelliJ.blank_lines(), self._stop_after).visit(tree, p, self._cursor.fork()) - tree = WrappingAndBracesVisitor(WrappingAndBracesStyle(), self._stop_after).visit(tree, p, self._cursor.fork()) - tree = SpacesVisitor(cu.get_style(SpacesStyle) or IntelliJ.spaces(), self._stop_after).visit(tree, p, self._cursor.fork()) tree = NormalizeTabsOrSpacesVisitor( diff --git a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py index 2d7f73a0..dfa7303f 100644 --- a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py +++ b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py @@ -168,7 +168,8 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], elif loc in (JRightPadded.Location.NEW_CLASS_ARGUMENTS, JRightPadded.Location.ARRAY_INDEX, JRightPadded.Location.PARENTHESES, - JRightPadded.Location.TYPE_PARAMETER): + JRightPadded.Location.TYPE_PARAMETER, + PyRightPadded.Location.COLLECTION_LITERAL_ELEMENT): # NOTE: DONE elem = self.visit_and_cast(elem, J, p) after = self._indent_to(right.after, indent, loc.after_location) @@ -223,9 +224,9 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], # NOTE: INCOMPLETE def visit_container(self, container: Optional[JContainer[J2]], loc: Union[PyContainer.Location, JContainer.Location], p: P) -> JContainer[J2]: - self._cursor = Cursor(self._cursor, container) if container is None: return container + self._cursor = Cursor(self._cursor, container) indent = cast(int, self.cursor.get_nearest_message("last_indent")) or 0 if '\n' in container.before.last_whitespace: From 0c9bc4095316367c260836eb79c67d53042cea6e Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Wed, 8 Jan 2025 17:08:11 +0100 Subject: [PATCH 11/16] Add fixed and some extra tests --- .../python/format/tabs_and_indents_visitor.py | 14 ++-- .../format/tabs_and_indents_visitor_test.py | 69 +++++++++++++++++-- 2 files changed, 72 insertions(+), 11 deletions(-) diff --git a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py index dfa7303f..fa4f6dff 100644 --- a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py +++ b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py @@ -147,7 +147,6 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], elif loc in (JRightPadded.Location.METHOD_DECLARATION_PARAMETER, JRightPadded.Location.RECORD_STATE_VECTOR): if isinstance(elem, Empty): - # NOTE: DONE elem = elem.with_prefix(self._indent_to(elem.prefix, indent, loc.after_location)) after = right.after else: @@ -163,16 +162,21 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], loc.after_location) elif loc == JRightPadded.Location.METHOD_INVOCATION_ARGUMENT: - # NOTE: DONE elem, after = self._visit_method_invocation_argument_j_type(elem, right, indent, loc, p) elif loc in (JRightPadded.Location.NEW_CLASS_ARGUMENTS, JRightPadded.Location.ARRAY_INDEX, JRightPadded.Location.PARENTHESES, - JRightPadded.Location.TYPE_PARAMETER, - PyRightPadded.Location.COLLECTION_LITERAL_ELEMENT): - # NOTE: DONE + JRightPadded.Location.TYPE_PARAMETER): elem = self.visit_and_cast(elem, J, p) after = self._indent_to(right.after, indent, loc.after_location) + elif loc in ( + PyRightPadded.Location.COLLECTION_LITERAL_ELEMENT, PyRightPadded.Location.DICT_LITERAL_ELEMENT): + elem = self.visit_and_cast(elem, J, p) + args = cast(JContainer[J], self.cursor.parent_or_throw.value) + # TODO: Maybe need to handle trailing comma? + if args.padding.elements[-1] is right: + self.cursor.parent_or_throw.put_message("indent_type", self.IndentType.ALIGN) + after = self.visit_space(right.after, loc.after_location, p) elif loc == JRightPadded.Location.ANNOTATION_ARGUMENT: raise NotImplementedError("Annotation argument not implemented") else: 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 ebe7faba..23a91789 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 @@ -375,7 +375,7 @@ def test_multiline_list(): 2, 3, 4 - ] + ] """ ), spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) @@ -485,6 +485,35 @@ def create_list(): spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) ) +def test_multiline_list_inside_function_with_trailing_comma(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) + rewrite_run( + # language=python + python( + """ + def create_list(): + my_list = [ + 1, + 2, + 3, + 4, + ] + return my_list + """, + """ + def create_list(): + my_list = [ + 1, + 2, + 3, + 4, + ] + return my_list + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + def test_basic_dictionary(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4).with_indent_size(4) @@ -509,6 +538,36 @@ def test_basic_dictionary(): def test_nested_dictionary(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + """ + config = { + "section": { + "key1": "value1", + "key2": [10, 20, + 30] + }, + "another_section": {"nested_key": "val"} + } + """, + """ + config = { + "section": { + "key1": "value1", + "key2": [10, 20, + 30] + }, + "another_section": {"nested_key": "val"} + } + """ + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_nested_dictionary_with_trailing_commas(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) rewrite_run( # language=python @@ -545,15 +604,13 @@ def test_list_comprehension(): python( """ def even_numbers(n): - return [ x for x in range(n) + return [x for x in range(n) if x % 2 == 0] """, """ def even_numbers(n): - return [ - x for x in range(n) - if x % 2 == 0 - ] + return [x for x in range(n) + if x % 2 == 0] """ ), spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) From e2b5dd27a07d338335ddb2cf482a9f3d4abb9798 Mon Sep 17 00:00:00 2001 From: Knut Wannheden Date: Thu, 9 Jan 2025 10:25:19 +0100 Subject: [PATCH 12/16] Add some trailing comma handling in collection literals --- .../python/format/tabs_and_indents_visitor.py | 51 +++++++++---------- .../format/tabs_and_indents_visitor_test.py | 8 +-- 2 files changed, 29 insertions(+), 30 deletions(-) diff --git a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py index fa4f6dff..d7391af4 100644 --- a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py +++ b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py @@ -7,7 +7,7 @@ from rewrite import Tree, Cursor, list_map from rewrite.java import J, Space, JRightPadded, JLeftPadded, JContainer, JavaSourceFile, EnumValueSet, Case, WhileLoop, \ Block, If, ForLoop, ForEachLoop, Package, Import, Label, DoWhileLoop, ArrayDimension, ClassDeclaration, Empty, \ - Binary, MethodInvocation, FieldAccess, Identifier, Lambda, TextComment, Comment + Binary, MethodInvocation, FieldAccess, Identifier, Lambda, TextComment, Comment, TrailingComma from rewrite.python import PythonVisitor, TabsAndIndentsStyle, PySpace, PyContainer, PyRightPadded, DictLiteral, \ CollectionLiteral from rewrite.visitor import P, T @@ -15,18 +15,18 @@ J2 = TypeVar('J2', bound=J) -class TabsAndIndentsVisitor(PythonVisitor): +class TabsAndIndentsVisitor(PythonVisitor[P]): - def __init__(self, style: TabsAndIndentsStyle, stop_after: Tree = None): + def __init__(self, style: TabsAndIndentsStyle, stop_after: Optional[Tree] = None): self._stop_after = stop_after self._style = style self._stop = False - def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> Optional[T]: + def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> Optional[J]: if parent is not None: self._cursor = parent if tree is None: - return cast(Optional[T], self.default_value(None, p)) + return cast(Optional[J], self.default_value(None, p)) for c in parent.get_path_as_cursors() if parent is not None else []: v = c.value @@ -74,13 +74,13 @@ def post_visit(self, tree: T, p: P) -> Optional[T]: def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Location, Space.Location]], p: P) -> Space: if space is None: - return space # pyright: ignore [reportReturnType] + return space # type: ignore self._cursor.put_message("last_location", loc) parent = self._cursor.parent indent = cast(int, self.cursor.get_nearest_message("last_indent")) or 0 - indent_type = self.cursor.parent.get_nearest_message("indent_type") or self.IndentType.ALIGN + indent_type = self.cursor.parent_or_throw.get_nearest_message("indent_type") or self.IndentType.ALIGN if not space.comments and '\n' not in space.last_whitespace or parent is None: return space @@ -89,7 +89,7 @@ def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Locati # Block spaces are always aligned to their parent # The second condition ensure init blocks are ignored. - # TODO: Secon condition might be removed since it's not relevant for Python + # TODO: Second condition might be removed since it's not relevant for Python align_block_prefix_to_parent = loc is Space.Location.BLOCK_PREFIX and '\n' in space.whitespace and \ (isinstance(cursor_value, Block) and not isinstance( self.cursor.parent_tree_cursor().value, Block)) @@ -102,7 +102,7 @@ def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Locati ) if (loc == Space.Location.EXTENDS and "\n" in space.whitespace) or \ - Space.Location.EXTENDS == self.cursor.parent.get_message("last_location", None): + Space.Location.EXTENDS == self.cursor.parent_or_throw.get_message("last_location", None): indent_type = self.IndentType.CONTINUATION_INDENT if align_block_prefix_to_parent or align_block_to_parent: @@ -140,6 +140,7 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], if isinstance(t, J): elem = t + trailing_comma = right.markers.find_first(TrailingComma) if '\n' in right.after.last_whitespace or '\n' in elem.prefix.last_whitespace: if loc in (JRightPadded.Location.FOR_CONDITION, JRightPadded.Location.FOR_UPDATE): @@ -169,27 +170,26 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], JRightPadded.Location.TYPE_PARAMETER): elem = self.visit_and_cast(elem, J, p) after = self._indent_to(right.after, indent, loc.after_location) - elif loc in ( - PyRightPadded.Location.COLLECTION_LITERAL_ELEMENT, PyRightPadded.Location.DICT_LITERAL_ELEMENT): + elif loc in (PyRightPadded.Location.COLLECTION_LITERAL_ELEMENT, PyRightPadded.Location.DICT_LITERAL_ELEMENT): elem = self.visit_and_cast(elem, J, p) args = cast(JContainer[J], self.cursor.parent_or_throw.value) - # TODO: Maybe need to handle trailing comma? - if args.padding.elements[-1] is right: + if not trailing_comma and args.padding.elements[-1] is right: self.cursor.parent_or_throw.put_message("indent_type", self.IndentType.ALIGN) after = self.visit_space(right.after, loc.after_location, p) + if trailing_comma: + self.cursor.parent_or_throw.put_message("indent_type", self.IndentType.ALIGN) + trailing_comma = trailing_comma.with_suffix(self.visit_space(trailing_comma.suffix, loc.after_location, p)) + right = right.with_markers(right.markers.compute_by_type(TrailingComma, lambda t: trailing_comma)) elif loc == JRightPadded.Location.ANNOTATION_ARGUMENT: raise NotImplementedError("Annotation argument not implemented") else: - # NOTE: DONE elem = self.visit_and_cast(elem, J, p) after = self.visit_space(right.after, loc.after_location, p) else: - # NOTE: Done if loc in (JRightPadded.Location.NEW_CLASS_ARGUMENTS, JRightPadded.Location.METHOD_INVOCATION_ARGUMENT): any_other_arg_on_own_line = False if "\n" not in elem.prefix.last_whitespace: - # NOTE: Done - args: JContainer[J] = cast(JContainer[J], self.cursor.parent.value) + args = cast(JContainer[J], self.cursor.parent_or_throw.value) for arg in args.padding.elements: if arg == self.cursor.value: continue @@ -201,7 +201,6 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], after = self._indent_to(right.after, indent, loc.after_location) if not any_other_arg_on_own_line: - # NOTE: Done if not isinstance(elem, Binary): if not isinstance(elem, MethodInvocation) or "\n" in elem.prefix.last_whitespace: self.cursor.put_message("last_indent", indent + self._style.continuation_indent) @@ -214,7 +213,6 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], elem = self.visit_and_cast(elem, J, p) after = self.visit_space(right.after, loc.after_location, p) else: - # NOTE: DONE elem = self.visit_and_cast(elem, J, p) after = self.visit_space(right.after, loc.after_location, p) @@ -222,14 +220,15 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], else: after = self.visit_space(right.after, loc.after_location, p) - self.cursor = self.cursor.parent + self.cursor = self.cursor.parent # type: ignore return right.with_after(after).with_element(t) # NOTE: INCOMPLETE def visit_container(self, container: Optional[JContainer[J2]], loc: Union[PyContainer.Location, JContainer.Location], p: P) -> JContainer[J2]: if container is None: - return container + return container # type: ignore + self._cursor = Cursor(self._cursor, container) indent = cast(int, self.cursor.get_nearest_message("last_indent")) or 0 @@ -257,14 +256,14 @@ def visit_container(self, container: Optional[JContainer[J2]], before = self.visit_space(container.before, loc.before_location, p) js = list_map(lambda t: self.visit_right_padded(t, loc.element_location, p), container.padding.elements) - self._cursor = self._cursor.parent + self._cursor = self._cursor.parent # type: ignore if container.padding.elements is js and container.before is before: return container return JContainer(before, js, container.markers) # NOTE: INCOMPLETE, Comments not supported - def _indent_to(self, space: Space, column: int, space_location: Space.Location) -> Space: + def _indent_to(self, space: Space, column: int, space_location: Union[PySpace.Location, Space.Location]) -> Space: s = space whitespace = s.whitespace @@ -281,7 +280,7 @@ def _indent_to(self, space: Space, column: int, space_location: Space.Location) else: def whitespace_indent(text: str) -> str: # TODO: Placeholder function, taken from java openrewrite.StringUtils - indent = [] + indent: List[str] = [] for c in text: if c == '\n' or c == '\r': return ''.join(indent) @@ -340,7 +339,7 @@ def _process_comment(i: int, c: Comment) -> Comment: s = s.with_comments(list_map(lambda i, c: _process_comment(c, i), s.comments)) return s - def _indent(self, whitespace: str, shift: int): + def _indent(self, whitespace: str, shift: int) -> str: return self._shift(whitespace, shift) def _indent_comment(self, comment: Comment, prior_suffix: str, to_column: int) -> Comment: @@ -378,7 +377,7 @@ def _get_length_of_whitespace(self, whitespace: Optional[str]) -> int: length = 0 return length - def _visit_method_invocation_argument_j_type(self, elem: J, right, indent, loc, p) -> tuple[J, Space]: + def _visit_method_invocation_argument_j_type(self, elem: J, right: JRightPadded[T], indent: int, loc: Union[PyRightPadded.Location, JRightPadded.Location], p: P) -> tuple[J, Space]: if "\n" not in elem.prefix.last_whitespace and isinstance(elem, Lambda): body = elem.body if not isinstance(body, Binary): 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 23a91789..3e4b7af3 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 @@ -557,7 +557,7 @@ def test_nested_dictionary(): "section": { "key1": "value1", "key2": [10, 20, - 30] + 30] }, "another_section": {"nested_key": "val"} } @@ -587,7 +587,7 @@ def test_nested_dictionary_with_trailing_commas(): "section": { "key1": "value1", "key2": [10, 20, - 30], + 30], }, "another_section": {"nested_key": "val"} } @@ -645,10 +645,10 @@ def my_function(a, b): if a > b: # cool a = b + 1 - # cool + # cool return None # Informative comment 4 - # Informative comment 5 + # Informative comment 5 ''' ), spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) From dfc85b164d06750a220468f3b14bf9ba97f97010 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Thu, 9 Jan 2025 11:14:54 +0100 Subject: [PATCH 13/16] Add extra test and some support code for METHOD_SELECT_SUFFIX --- .../python/format/tabs_and_indents_visitor.py | 8 +++- .../format/tabs_and_indents_visitor_test.py | 42 +++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py index d7391af4..1a10ce32 100644 --- a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py +++ b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py @@ -79,6 +79,12 @@ def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Locati self._cursor.put_message("last_location", loc) parent = self._cursor.parent + if loc == Space.Location.METHOD_SELECT_SUFFIX: + chained_indent = self.cursor.parent_tree_cursor().get_message("chained_indent", None) + if chained_indent is not None: + self.cursor.parent_tree_cursor().put_message("last_indent", chained_indent) + return self._indent_to(space, chained_indent, loc) + indent = cast(int, self.cursor.get_nearest_message("last_indent")) or 0 indent_type = self.cursor.parent_or_throw.get_nearest_message("indent_type") or self.IndentType.ALIGN @@ -117,7 +123,7 @@ def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Locati if isinstance(cursor_value, J) and not isinstance(cursor_value, EnumValueSet): self.cursor.put_message("last_indent", indent) elif loc == Space.Location.METHOD_SELECT_SUFFIX: - raise NotImplementedError("Method select suffix not implemented") + self.cursor.parent_tree_cursor().put_message("last_indent", indent) return s 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 3e4b7af3..a3535354 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 @@ -680,3 +680,45 @@ def my_function(): ), spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) ) + + +def test_method_select_suffix(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + """ + x = ("foo" + .startswith("f")) + """, + """ + x = ("foo" + .startswith("f")) + """), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_method_select_suffix_new_line_already_correct(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + """ + x = ("foo" + .startswith("f")) + """), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + + +def test_method_select_suffix_already_correct(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + """ + x = ("foo".startswith("f")) + """), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) From 9e1c596a1de135270e7143a707f9f50e76f97a56 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Thu, 9 Jan 2025 11:18:11 +0100 Subject: [PATCH 14/16] Reorder --- .../format/tabs_and_indents_visitor_test.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) 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 a3535354..95996577 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 @@ -682,18 +682,13 @@ def my_function(): ) -def test_method_select_suffix(): +def test_method_select_suffix_already_correct(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) rewrite_run( # language=python python( """ - x = ("foo" - .startswith("f")) - """, - """ - x = ("foo" - .startswith("f")) + x = ("foo".startswith("f")) """), spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) ) @@ -712,13 +707,18 @@ def test_method_select_suffix_new_line_already_correct(): ) -def test_method_select_suffix_already_correct(): +def test_method_select_suffix(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) rewrite_run( # language=python python( """ - x = ("foo".startswith("f")) + x = ("foo" + .startswith("f")) + """, + """ + x = ("foo" + .startswith("f")) """), spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) ) From cd7e7963affed51d83698e5d9f72f1dcb9867503 Mon Sep 17 00:00:00 2001 From: Niels de Bruin Date: Thu, 9 Jan 2025 11:56:02 +0100 Subject: [PATCH 15/16] Add extra test cases for comments and mark some test with xfail --- .../format/tabs_and_indents_visitor_test.py | 35 ++++++++++++++++--- 1 file changed, 31 insertions(+), 4 deletions(-) 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 95996577..28931c2d 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 @@ -1,3 +1,5 @@ +import pytest + from rewrite.python import IntelliJ from rewrite.python.format import TabsAndIndentsVisitor from rewrite.test import rewrite_run, python, RecipeSpec, from_visitor @@ -631,8 +633,6 @@ def my_function(a, b): if a > b: # cool a = b + 1 - # cool - return None # Informative comment 4 # Informative comment 5 ''', @@ -645,8 +645,6 @@ def my_function(a, b): if a > b: # cool a = b + 1 - # cool - return None # Informative comment 4 # Informative comment 5 ''' @@ -656,6 +654,33 @@ def my_function(a, b): return None +def test_comment_alignment_if_and_return(): + style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) + rewrite_run( + # language=python + python( + ''' + def my_function(a, b): + if a > b: + # cool + a = b + 1 + # cool + return None + ''', + ''' + def my_function(a, b): + if a > b: + # cool + a = b + 1 + # cool + return None + ''' + ), + spec=RecipeSpec().with_recipes(from_visitor(TabsAndIndentsVisitor(style))) + ) + return None + + def test_docstring_alignment(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) rewrite_run( @@ -694,6 +719,7 @@ def test_method_select_suffix_already_correct(): ) +@pytest.mark.xfail def test_method_select_suffix_new_line_already_correct(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) rewrite_run( @@ -707,6 +733,7 @@ def test_method_select_suffix_new_line_already_correct(): ) +@pytest.mark.xfail def test_method_select_suffix(): style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4) rewrite_run( From ae81f3c195a3bf22e6657f8d1762ec9b90cac3ce Mon Sep 17 00:00:00 2001 From: Knut Wannheden Date: Thu, 9 Jan 2025 11:57:55 +0100 Subject: [PATCH 16/16] Polish `TabsAndIndentsVisitor` --- .../python/format/tabs_and_indents_visitor.py | 37 +++++++------------ .../format/tabs_and_indents_visitor_test.py | 2 +- 2 files changed, 15 insertions(+), 24 deletions(-) diff --git a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py index 1a10ce32..b486a109 100644 --- a/rewrite/rewrite/python/format/tabs_and_indents_visitor.py +++ b/rewrite/rewrite/python/format/tabs_and_indents_visitor.py @@ -5,11 +5,11 @@ from typing import TypeVar, Optional, Union, cast, List from rewrite import Tree, Cursor, list_map -from rewrite.java import J, Space, JRightPadded, JLeftPadded, JContainer, JavaSourceFile, EnumValueSet, Case, WhileLoop, \ - Block, If, ForLoop, ForEachLoop, Package, Import, Label, DoWhileLoop, ArrayDimension, ClassDeclaration, Empty, \ +from rewrite.java import J, Space, JRightPadded, JLeftPadded, JContainer, JavaSourceFile, Case, WhileLoop, \ + Block, If, Label, ArrayDimension, ClassDeclaration, Empty, \ Binary, MethodInvocation, FieldAccess, Identifier, Lambda, TextComment, Comment, TrailingComma from rewrite.python import PythonVisitor, TabsAndIndentsStyle, PySpace, PyContainer, PyRightPadded, DictLiteral, \ - CollectionLiteral + CollectionLiteral, ForLoop from rewrite.visitor import P, T J2 = TypeVar('J2', bound=J) @@ -28,7 +28,7 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> if tree is None: return cast(Optional[J], self.default_value(None, p)) - for c in parent.get_path_as_cursors() if parent is not None else []: + for c in parent.get_path_as_cursors(): v = c.value space = None if isinstance(v, J): @@ -45,20 +45,18 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> if indent != 0: c.put_message("last_indent", indent) - # TODO CHeck if this is 100% percent corrent - for v in parent.get_path() if parent is not None else []: - if isinstance(v, J): - self.pre_visit(v, p) + for next_parent in parent.get_path(): + if isinstance(next_parent, J): + self.pre_visit(next_parent, p) break + return super().visit(tree, p) def pre_visit(self, tree: T, p: P) -> Optional[T]: - if isinstance(tree, ( - JavaSourceFile, Package, Import, Label, DoWhileLoop, ArrayDimension, ClassDeclaration)): + if isinstance(tree, (JavaSourceFile, Label, ArrayDimension, ClassDeclaration)): self.cursor.put_message("indent_type", self.IndentType.ALIGN) elif isinstance(tree, - (Block, If, If.Else, ForLoop, ForEachLoop, WhileLoop, Case, EnumValueSet, DictLiteral, - CollectionLiteral)): + (Block, If, If.Else, ForLoop, WhileLoop, Case, DictLiteral, CollectionLiteral)): # NOTE: Added CollectionLiteral, DictLiteral here self.cursor.put_message("indent_type", self.IndentType.INDENT) else: @@ -80,7 +78,7 @@ def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Locati parent = self._cursor.parent if loc == Space.Location.METHOD_SELECT_SUFFIX: - chained_indent = self.cursor.parent_tree_cursor().get_message("chained_indent", None) + chained_indent = cast(int, self.cursor.parent_tree_cursor().get_message("chained_indent", None)) if chained_indent is not None: self.cursor.parent_tree_cursor().put_message("last_indent", chained_indent) return self._indent_to(space, chained_indent, loc) @@ -120,7 +118,7 @@ def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Locati indent += self._style.continuation_indent s: Space = self._indent_to(space, indent, loc) - if isinstance(cursor_value, J) and not isinstance(cursor_value, EnumValueSet): + if isinstance(cursor_value, J): self.cursor.put_message("last_indent", indent) elif loc == Space.Location.METHOD_SELECT_SUFFIX: self.cursor.parent_tree_cursor().put_message("last_indent", indent) @@ -128,7 +126,6 @@ def visit_space(self, space: Optional[Space], loc: Optional[Union[PySpace.Locati return s - # NOTE: INCOMPLETE def visit_right_padded(self, right: Optional[JRightPadded[T]], loc: Union[PyRightPadded.Location, JRightPadded.Location], p: P) -> Optional[ JRightPadded[T]]: @@ -229,7 +226,6 @@ def visit_right_padded(self, right: Optional[JRightPadded[T]], self.cursor = self.cursor.parent # type: ignore return right.with_after(after).with_element(t) - # NOTE: INCOMPLETE def visit_container(self, container: Optional[JContainer[J2]], loc: Union[PyContainer.Location, JContainer.Location], p: P) -> JContainer[J2]: if container is None: @@ -268,8 +264,7 @@ def visit_container(self, container: Optional[JContainer[J2]], return container return JContainer(before, js, container.markers) - # NOTE: INCOMPLETE, Comments not supported - def _indent_to(self, space: Space, column: int, space_location: Union[PySpace.Location, Space.Location]) -> Space: + def _indent_to(self, space: Space, column: int, space_location: Optional[Union[PySpace.Location, Space.Location]]) -> Space: s = space whitespace = s.whitespace @@ -333,7 +328,7 @@ def _process_comment(i: int, c: Comment) -> Comment: new_c = c if "\n" in prior_suffix or has_file_leading_comment: - new_c = self._indent_comment(c, prior_suffix, to_column) + new_c = c if '\n' in new_c.suffix: suffix_indent = self._get_length_of_whitespace(new_c.suffix) @@ -348,10 +343,6 @@ def _process_comment(i: int, c: Comment) -> Comment: def _indent(self, whitespace: str, shift: int) -> str: return self._shift(whitespace, shift) - def _indent_comment(self, comment: Comment, prior_suffix: str, to_column: int) -> Comment: - # TODO: Handle multiline here. - return comment - def _shift(self, text: str, shift: int) -> str: tab_indent = self._style.tab_size if not self._style.use_tab_character: 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 28931c2d..6b8e78b3 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 @@ -672,7 +672,7 @@ def my_function(a, b): if a > b: # cool a = b + 1 - # cool + # cool return None ''' ),