Support // and ** operators

openrewrite · Sep 5, 2024 · f1370c9 · f1370c9
1 parent efb7a74
commit f1370c9
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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
@@ -115,6 +115,12 @@ public J visitBinary(Py.Binary binary, PrintOutputCapture<P> 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
@@ -127,6 +127,9 @@ public enum Type {
             Is,
             IsNot,
             NotIn,
+
+            FloorDivision,
+            Power
         }
 
         public Padding getPadding() {
-Original file line number
+Diff line change
@@ Expand Up / @@ -127,6 +127,9 @@ public enum Type { @@
                 Is,
                 IsNot,
                 NotIn,
+                FloorDivision,
+                Power
             }
             public Padding getPadding() {
@@ Expand Down @@