diff --git a/rewrite/rewrite/python/_parser_visitor.py b/rewrite/rewrite/python/_parser_visitor.py index 2c3bf924..f35a4315 100644 --- a/rewrite/rewrite/python/_parser_visitor.py +++ b/rewrite/rewrite/python/_parser_visitor.py @@ -987,15 +987,31 @@ def visit_Ellipsis(self, node): def visit_BinOp(self, node): - return j.Binary( - random_id(), - self.__whitespace(), - Markers.EMPTY, - self.__convert(node.left), - self.__convert_binary_operator(node.op), - self.__convert(node.right), - self.__map_type(node) - ) + prefix = self.__whitespace() + left = self.__convert(node.left) + op = self.__convert_binary_operator(node.op) + + if isinstance(op.element, py.Binary.Type): + return py.Binary( + random_id(), + prefix, + Markers.EMPTY, + left, + op, + None, + self.__convert(node.right), + self.__map_type(node) + ) + else: + return j.Binary( + random_id(), + prefix, + Markers.EMPTY, + left, + op, + self.__convert(node.right), + self.__map_type(node) + ) def visit_BoolOp(self, node): @@ -1110,6 +1126,7 @@ def __convert_binary_operator(self, op) -> Union[JLeftPadded[j.Binary.Type], JLe ast.BitXor: (j.Binary.Type.BitXor, '^'), ast.Div: (j.Binary.Type.Division, '/'), ast.Eq: (j.Binary.Type.Equal, '=='), + ast.FloorDiv: (py.Binary.Type.FloorDivision, '//'), ast.Gt: (j.Binary.Type.GreaterThan, '>'), ast.GtE: (j.Binary.Type.GreaterThanOrEqual, '>='), ast.In: (py.Binary.Type.In, 'in'), @@ -1123,6 +1140,7 @@ def __convert_binary_operator(self, op) -> Union[JLeftPadded[j.Binary.Type], JLe ast.NotEq: (j.Binary.Type.NotEqual, '!='), ast.NotIn: (py.Binary.Type.NotIn, 'not'), ast.Or: (j.Binary.Type.Or, 'or'), + ast.Pow: (py.Binary.Type.Power, '**'), ast.RShift: (j.Binary.Type.RightShift, '>>'), ast.Sub: (j.Binary.Type.Subtraction, '-'), } @@ -1716,12 +1734,24 @@ def __pad_left(self, space: Space, tree) -> JLeftPadded[J2]: def __source_before(self, until_delim: str, stop: Optional[str] = None) -> Space: + if self._source.startswith(until_delim, self._cursor): + self._cursor += len(until_delim) + return Space.EMPTY + + save_cursor = self._cursor + space = self.__whitespace() + if self._source.startswith(until_delim, self._cursor): + self._cursor += len(until_delim) + return space + else: + self._cursor = save_cursor + delim_index = self.__position_of_next(until_delim, stop) if delim_index == -1: return Space.EMPTY if delim_index == self._cursor: - self._cursor = self._cursor + len(until_delim) + self._cursor += len(until_delim) return Space.EMPTY space = self.__whitespace() @@ -1782,7 +1812,8 @@ def __position_of_next(self, until_delim: str, stop: str = None) -> int: in_single_line_comment = False delim_index = self._cursor - while delim_index < len(self._source) - len(until_delim) + 1: + end_index = len(self._source) - len(until_delim) + 1 + while delim_index < end_index: if in_single_line_comment: if self._source[delim_index] == '\n': in_single_line_comment = False diff --git a/rewrite/rewrite/python/tree.py b/rewrite/rewrite/python/tree.py index d629a31c..ebfe278c 100644 --- a/rewrite/rewrite/python/tree.py +++ b/rewrite/rewrite/python/tree.py @@ -94,6 +94,8 @@ class Type(Enum): Is = 1 IsNot = 2 NotIn = 3 + FloorDivision = 4 + Power = 5 @dataclass class PaddingHelper: diff --git a/rewrite/tests/python/all/binary_test.py b/rewrite/tests/python/all/binary_test.py index 3028fda0..88e9a6ca 100644 --- a/rewrite/tests/python/all/binary_test.py +++ b/rewrite/tests/python/all/binary_test.py @@ -21,6 +21,10 @@ def test_arithmetic_ops(): rewrite_run(python("assert 1 / 2")) # language=python rewrite_run(python("assert 1 % 2")) + # language=python + rewrite_run(python("assert 1 // 2")) + # language=python + rewrite_run(python("assert 1 ** 2")) def test_eq_ops(): diff --git a/src/main/java/org/openrewrite/python/internal/PythonPrinter.java b/src/main/java/org/openrewrite/python/internal/PythonPrinter.java index 939ac23c..90cbfff4 100755 --- a/src/main/java/org/openrewrite/python/internal/PythonPrinter.java +++ b/src/main/java/org/openrewrite/python/internal/PythonPrinter.java @@ -115,6 +115,12 @@ public J visitBinary(Py.Binary binary, PrintOutputCapture

p) { } p.append("not"); break; + case FloorDivision: + p.append("//"); + break; + case Power: + p.append("**"); + break; } visit(binary.getRight(), p); diff --git a/src/main/java/org/openrewrite/python/tree/Py.java b/src/main/java/org/openrewrite/python/tree/Py.java index a1bd85e5..1ced2511 100644 --- a/src/main/java/org/openrewrite/python/tree/Py.java +++ b/src/main/java/org/openrewrite/python/tree/Py.java @@ -127,6 +127,9 @@ public enum Type { Is, IsNot, NotIn, + + FloorDivision, + Power } public Padding getPadding() {