openrewrite-remote 0.13.4__py3-none-any.whl → 0.14.0__py3-none-any.whl

rewrite_remote/remote_utils.py

@@ -57,9 +57,7 @@ def receive_nodes(
             elif diff_event.event_type == EventType.NoChange:
                 after[i] = None  # Or some default value
             else:
-                raise NotImplementedError(
-                    f"Unexpected operation: {diff_event.event_type}"
-                )
+                raise NotImplementedError(f"Unexpected operation: {diff_event.event_type}")
         return after  # type: ignore
     elif list_event.event_type == EventType.Update:
         return _receive_updated_nodes(before, list_event.msg, details, ctx)  # type: ignore
@@ -85,7 +83,11 @@ def _receive_updated_nodes(
         if evt.event_type in (EventType.NoChange, EventType.EndList):
             break
 
-        if evt.event_type in (EventType.Delete, EventType.Update, EventType.Add):
+        if evt.event_type in (
+            EventType.Delete,
+            EventType.Update,
+            EventType.Add,
+        ):
             if not modified:
                 after_list = _copy_range(before, before_idx)
                 modified = True
@@ -97,9 +99,7 @@ def _receive_updated_nodes(
         elif evt.event_type == EventType.Delete:
             before_idx += 1
         elif evt.event_type == EventType.Update:
-            after_list.append(
-                details.receive_details(before[before_idx], evt.concrete_type, ctx)
-            )
+            after_list.append(details.receive_details(before[before_idx], evt.concrete_type, ctx))
             before_idx += 1
         elif evt.event_type == EventType.Add:
             after_list.append(details.receive_details(None, evt.concrete_type, ctx))
@@ -130,9 +130,7 @@ def receive_values(
             elif diff_event.event_type == EventType.NoChange:
                 after[i] = None  # Or some default value
             else:
-                raise NotImplementedError(
-                    f"Unexpected operation: {diff_event.event_type}"
-                )
+                raise NotImplementedError(f"Unexpected operation: {diff_event.event_type}")
         return after  # type: ignore
     elif list_event.event_type == EventType.Update:
         return _receive_updated_values(before, list_event.msg, type, ctx)  # type: ignore
@@ -155,7 +153,11 @@ def _receive_updated_values(
         if evt.event_type in (EventType.NoChange, EventType.EndList):
             break
 
-        if evt.event_type in (EventType.Delete, EventType.Update, EventType.Add):
+        if evt.event_type in (
+            EventType.Delete,
+            EventType.Update,
+            EventType.Add,
+        ):
             if not modified:
                 after_list = _copy_range(before, before_idx)
                 modified = True
@@ -239,16 +241,20 @@ def calculate_list_diff(
 
                 # If elements at current indices are not equal, figure out the operation
                 if before_id not in after_map:
-                    consumer(Operation.Delete, before_idx, -1, before[before_idx], None)
+                    consumer(
+                        Operation.Delete,
+                        before_idx,
+                        -1,
+                        before[before_idx],
+                        None,
+                    )
                     before_idx += 1
                 else:
                     consumer(Operation.Add, -1, after_idx, None, after[after_idx])
                     after_idx += 1
 
 
-def create_index_map(
-    lst: List[T], from_index: int, id_function: Callable[[T], I]
-) -> Dict[I, int]:
+def create_index_map(lst: List[T], from_index: int, id_function: Callable[[T], I]) -> Dict[I, int]:
     result = {}
     for i in range(from_index, len(lst)):
         result[id_function(lst[i])] = i

rewrite_remote/remoting.py

@@ -12,8 +12,6 @@ from typing import (
     Dict,
     Optional,
     Type,
-    List,
-    Tuple,
     Callable,
     cast,
     Union,
@@ -108,14 +106,16 @@ class RemotingContext:
     def new_sender_context(self, output_stream: Any) -> "SenderContext":
         return SenderContext(
             JsonSender(
-                output_stream, SerializationContext(self, self._value_serializers)
+                output_stream,
+                SerializationContext(self, self._value_serializers),
             )
         )
 
     def new_receiver_context(self, input_stream: Any) -> "ReceiverContext":
         return ReceiverContext(
             JsonReceiver(
-                input_stream, DeserializationContext(self, self._value_deserializers)
+                input_stream,
+                DeserializationContext(self, self._value_deserializers),
             )
         )
 
@@ -160,7 +160,6 @@ ERROR = 1
 
 
 class RemotingMessenger:
-
     def __init__(
         self,
         context: RemotingContext,
@@ -249,9 +248,7 @@ class RemotingMessenger:
         root_cursor = Cursor(None, Cursor.ROOT_VALUE)
         ctx = InMemoryExecutionContext()
         RemotingExecutionContextView.view(ctx).remoting_context = self._context
-        print_output = received.print(
-            Cursor(root_cursor, received), PrintOutputCapture(0)
-        )
+        print_output = received.print(Cursor(root_cursor, received), PrintOutputCapture(0))
 
         response_stream = BytesIO()
         cbor2.dump(RemotingMessageType.Response, response_stream)
@@ -280,7 +277,9 @@ class RemotingMessenger:
 
     def send_print_request(self, sock: socket.socket, cursor: Cursor):
         self.__send_request_stream(
-            sock, "print", lambda s: self.send_tree(s, cast(Tree, cursor.value))
+            sock,
+            "print",
+            lambda s: self.send_tree(s, cast(Tree, cursor.value)),
         )
         if self.recv_byte(sock) != RemotingMessageType.Response:
             raise ValueError("Unexpected message type.")
@@ -293,16 +292,22 @@ class RemotingMessenger:
 
     def send_tree(
         self,
-        sock: Union[BinaryIO, socket.socket],
+        dest: Union[BinaryIO, socket.socket],
         after: Tree,
         before: Optional[Tree] = None,
     ):
         b = BytesIO()
         self._context.new_sender_context(b).send_any_tree(after, before)
-        sock.sendall(dumps(b.getvalue()))
+        if isinstance(dest, socket.socket):
+            dest.sendall(dumps(b.getvalue()))
+        else:
+            b.seek(0)
+            dest.write(dumps(b.getvalue()))
 
     def receive_tree(
-        self, data: Union[BinaryIO, socket.socket], before: Optional[Tree] = None
+        self,
+        data: Union[BinaryIO, socket.socket],
+        before: Optional[Tree] = None,
     ):
         receiver_context = self._context.new_receiver_context(BytesIO(cbor2.load(data)))
         return receiver_context.receive_any_tree(before)

rewrite_remote/sender.py

@@ -21,7 +21,6 @@ import cbor2
 from cbor2 import CBOREncoder
 from rewrite import Tree, Markers, Marker, ParseErrorVisitor, Cursor, Style
 from rewrite.visitor import TreeVisitor
-from rewrite.java import JavaType
 
 from rewrite_remote import remote_utils
 from rewrite_remote.event import *
@@ -52,9 +51,7 @@ class OmniSender(Sender):
 
 
 class TreeSender(Protocol):
-    def send_node(
-            self, diff_event: DiffEvent, visitor: Callable[["TreeSender"], None]
-    ) -> None: ...
+    def send_node(self, diff_event: DiffEvent, visitor: Callable[["TreeSender"], None]) -> None: ...
 
     def send_value(self, diff_event: DiffEvent) -> None: ...
 
@@ -69,10 +66,10 @@ class SenderContext(Generic[T]):
     Registry: Dict[Type[Any], Callable[[], Sender]] = {}
 
     def __init__(
-            self,
-            sender: "TreeSender",
-            visitor: TreeVisitor = None,
-            before: Optional[Any] = None,
+        self,
+        sender: "TreeSender",
+        visitor: TreeVisitor = None,
+        before: Optional[Any] = None,
     ):
         self.sender = sender
         self.visitor = visitor
@@ -82,24 +79,20 @@ class SenderContext(Generic[T]):
         for entry_type, factory in self.Registry.items():
             # FIXME find better solution
             try:
-                if type_.__bases__.__contains__(entry_type) or issubclass(
-                        type_, entry_type
-                ):
+                if type_.__bases__.__contains__(entry_type) or issubclass(type_, entry_type):
                     return factory()
             except:
                 pass
         raise ValueError(f"Unsupported sender type: {type_}")
 
-    def fork(
-            self, visitor: TreeVisitor, before: Optional[Any] = None
-    ) -> "SenderContext[T]":
+    def fork(self, visitor: TreeVisitor, before: Optional[Any] = None) -> "SenderContext[T]":
         return SenderContext(self.sender, visitor, before)
 
     def visit(
-            self,
-            consumer: Callable[[V, "SenderContext"], None],
-            after: V,
-            before: Optional[V] = None,
+        self,
+        consumer: Callable[[V, "SenderContext"], None],
+        after: V,
+        before: Optional[V] = None,
     ) -> None:
         save_before = self.before
         self.before = before
@@ -113,10 +106,10 @@ class SenderContext(Generic[T]):
         OmniSender().send(after, before, self)
 
     def send_node(
-            self,
-            owner: A,
-            extractor: Callable[[A], Optional[V]],
-            details: Callable[[V, "SenderContext[T]"], None],
+        self,
+        owner: A,
+        extractor: Callable[[A], Optional[V]],
+        details: Callable[[V, "SenderContext[T]"], None],
     ) -> None:
         self.send_node_internal(
             extractor(owner),
@@ -124,9 +117,7 @@ class SenderContext(Generic[T]):
             details,
         )
 
-    def send_value(
-            self, owner: Any, value_extractor: Callable[[T], Optional[V]]
-    ) -> None:
+    def send_value(self, owner: Any, value_extractor: Callable[[T], Optional[V]]) -> None:
         after_value = value_extractor(owner)
         before_value = value_extractor(self.before) if self.before is not None else None
         self.send_value_internal(after_value, before_value)
@@ -136,9 +127,7 @@ class SenderContext(Generic[T]):
         before_value = value_extractor(self.before) if self.before is not None else None
         self.send_typed_value_internal(after_value, before_value)
 
-    def send_list_event(
-            self, after: Optional[List[V]], before: Optional[List[V]]
-    ) -> bool:
+    def send_list_event(self, after: Optional[List[V]], before: Optional[List[V]]) -> bool:
         if after == before:
             evt = DiffEvent(EventType.NoChange, None, None)
         elif before is None:
@@ -149,21 +138,17 @@ class SenderContext(Generic[T]):
             evt = DiffEvent(EventType.Update, None, len(after))
 
         self.sender.send_value(evt)
-        return (
-                evt.event_type != EventType.NoChange and evt.event_type != EventType.Delete
-        )
+        return evt.event_type != EventType.NoChange and evt.event_type != EventType.Delete
 
     def send_nodes(
-            self,
-            owner: A,
-            element_extractor: Callable[[A], List[V]],
-            details: Callable[[V, "SenderContext"], None],
-            id_function: Callable[[V], I],
+        self,
+        owner: A,
+        element_extractor: Callable[[A], List[V]],
+        details: Callable[[V, "SenderContext"], None],
+        id_function: Callable[[V], I],
     ) -> None:
         after_list = element_extractor(owner)
-        before_list = (
-            element_extractor(self.before) if self.before is not None else None
-        )
+        before_list = element_extractor(self.before) if self.before is not None else None
 
         if self.send_list_event(after_list, before_list):
             if before_list is not None:
@@ -175,14 +160,10 @@ class SenderContext(Generic[T]):
                 id_function,
                 lambda op, _1, _2, bv, av: {
                     Operation.Delete: lambda: self.send_node_internal(av, bv, details),
-                    Operation.NoChange: lambda: self.send_node_internal(
-                        av, bv, details
-                    ),
+                    Operation.NoChange: lambda: self.send_node_internal(av, bv, details),
                     Operation.Add: lambda: self.send_node_internal(av, bv, details),
                     Operation.Update: lambda: self.send_node_internal(av, bv, details),
-                    Operation.Move: lambda: NotImplementedError(
-                        "Unexpected operation: " + str(op)
-                    ),
+                    Operation.Move: lambda: NotImplementedError("Unexpected operation: " + str(op)),
                 }[op](),
             )
 
@@ -190,10 +171,10 @@ class SenderContext(Generic[T]):
                 self.sender.send_value(DiffEvent(EventType.EndList, None, None))
 
     def send_values(
-            self,
-            owner: T,
-            value_extractor: Callable[[T], List[V]],
-            id_function: Callable[[V], I],
+        self,
+        owner: T,
+        value_extractor: Callable[[T], List[V]],
+        id_function: Callable[[V], I],
     ) -> None:
         after_list = value_extractor(owner)
         before_list = value_extractor(self.before) if self.before is not None else None
@@ -211,9 +192,7 @@ class SenderContext(Generic[T]):
                     Operation.NoChange: lambda: self.send_value_internal(av, bv),
                     Operation.Add: lambda: self.send_value_internal(av, bv),
                     Operation.Update: lambda: self.send_value_internal(av, bv),
-                    Operation.Move: lambda: NotImplementedError(
-                        "Unexpected operation: " + str(op)
-                    ),
+                    Operation.Move: lambda: NotImplementedError("Unexpected operation: " + str(op)),
                 }[op](),
             )
 
@@ -221,10 +200,10 @@ class SenderContext(Generic[T]):
                 self.sender.send_value(DiffEvent(EventType.EndList, None, None))
 
     def send_typed_values(
-            self,
-            owner: T,
-            value_extractor: Callable[[T], List[V]],
-            id_function: Callable[[V], I],
+        self,
+        owner: T,
+        value_extractor: Callable[[T], List[V]],
+        id_function: Callable[[V], I],
     ) -> None:
         after_list = value_extractor(owner)
         before_list = value_extractor(self.before) if self.before is not None else None
@@ -242,9 +221,7 @@ class SenderContext(Generic[T]):
                     Operation.NoChange: lambda: self.send_typed_value_internal(av, bv),
                     Operation.Add: lambda: self.send_typed_value_internal(av, bv),
                     Operation.Update: lambda: self.send_typed_value_internal(av, bv),
-                    Operation.Move: lambda: NotImplementedError(
-                        "Unexpected operation: " + str(op)
-                    ),
+                    Operation.Move: lambda: NotImplementedError("Unexpected operation: " + str(op)),
                 }[op](),
             )
 
@@ -270,17 +247,15 @@ class SenderContext(Generic[T]):
         return isinstance(after, (Tree, Markers)) or isinstance(before, (Tree, Markers))
 
     def send_node_internal(
-            self,
-            after: Optional[V],
-            before: Optional[V],
-            details: Callable[[V, "SenderContext"], None],
+        self,
+        after: Optional[V],
+        before: Optional[V],
+        details: Callable[[V, "SenderContext"], None],
     ) -> None:
         if self.are_equal(after, before):
             evt = DiffEvent(EventType.NoChange, None, None)
         elif before is None:
-            concrete_type = (
-                to_java_type_name(type(after)) if after is not None else None
-            )
+            concrete_type = to_java_type_name(type(after)) if after is not None else None
             evt = DiffEvent(EventType.Add, concrete_type, None)
         elif after is None:
             evt = DiffEvent(EventType.Delete, None, None)
@@ -293,9 +268,7 @@ class SenderContext(Generic[T]):
         if self.before is not None and self.are_equal(after, before):
             evt = DiffEvent(EventType.NoChange, None, None)
         elif self.before is None or before is None:
-            concrete_type = (
-                to_java_type_name(type(after)) if isinstance(after, Marker) else None
-            )
+            concrete_type = to_java_type_name(type(after)) if isinstance(after, Marker) else None
             evt = DiffEvent(EventType.Add, concrete_type, after)
         elif after is None:
             evt = DiffEvent(EventType.Delete, None, None)
@@ -308,9 +281,7 @@ class SenderContext(Generic[T]):
         if self.before is not None and self.are_equal(after, before):
             evt = DiffEvent(EventType.NoChange, None, None)
         elif self.before is None or before is None:
-            concrete_type = (
-                to_java_type_name_from_value(after) if after is not None else None
-            )
+            concrete_type = to_java_type_name_from_value(after) if after is not None else None
             evt = DiffEvent(EventType.Add, concrete_type, after)
         elif after is None:
             evt = DiffEvent(EventType.Delete, None, None)
@@ -322,23 +293,21 @@ class SenderContext(Generic[T]):
 
 class SerializationContext:
     def __init__(
-            self,
-            remoting_context: "RemotingContext",
-            value_serializers: Optional[Dict[Type[Any], "ValueSerializer"]] = None,
+        self,
+        remoting_context: "RemotingContext",
+        value_serializers: Optional[Dict[Type[Any], "ValueSerializer"]] = None,
     ):
         self.remoting_context = remoting_context
         self.value_serializers = value_serializers or {}
 
-    def serialize(
-            self, value: Any, type_name: Optional[str], encoder: CBOREncoder
-    ) -> None:
+    def serialize(self, value: Any, type_name: Optional[str], encoder: CBOREncoder) -> None:
         if value is None:
             encoder.encode_none(None)
             return
 
         for type_cls, serializer in self.value_serializers.items():
             if isinstance(value, type_cls):
-                serializer.serialize(value, type_name, encoder, self)
+                serializer(value, type_name, encoder, self)
                 return
 
         DefaultValueSerializer().serialize(value, type_name, encoder, self)
@@ -353,7 +322,6 @@ class JsonSender(TreeSender):
         self._encoder = cbor2.CBOREncoder(self._stream)
 
     def send_node(self, diff_event, visitor) -> None:
-
         if diff_event.event_type in (EventType.Add, EventType.Update):
             self._encoder.encode(
                 [diff_event.event_type.value]
@@ -375,20 +343,16 @@ class JsonSender(TreeSender):
 
     def send_value(self, diff_event):
         if diff_event.event_type in (EventType.Add, EventType.Update):
-            self._encoder.encode_length(
-                4, 3 if diff_event.concrete_type is not None else 2
-            )
+            self._encoder.encode_length(4, 3 if diff_event.concrete_type is not None else 2)
             self._encoder.encode_int(diff_event.event_type.value)
             if diff_event.concrete_type is not None:
                 self._encoder.encode(diff_event.concrete_type)
-            self._context.serialize(
-                diff_event.msg, diff_event.concrete_type, self._encoder
-            )
+            self._context.serialize(diff_event.msg, diff_event.concrete_type, self._encoder)
         elif diff_event.event_type in (
-                EventType.Delete,
-                EventType.NoChange,
-                EventType.StartList,
-                EventType.EndList,
+            EventType.Delete,
+            EventType.NoChange,
+            EventType.StartList,
+            EventType.EndList,
         ):
             self._encoder.encode([diff_event.event_type.value])
         else:
@@ -409,11 +373,11 @@ ValueSerializer = Callable[[Any, Optional[str], CBOREncoder, SerializationContex
 
 
 def write_object_using_reflection(
-        value: Any,
-        type_name: Optional[str],
-        with_id: bool,
-        encoder: CBOREncoder,
-        context: SerializationContext,
+    value: Any,
+    type_name: Optional[str],
+    with_id: bool,
+    encoder: CBOREncoder,
+    context: SerializationContext,
 ) -> None:
     if with_id and (id := context.remoting_context.try_get_id(value)):
         encoder.encode_int(id)
@@ -429,8 +393,7 @@ def write_object_using_reflection(
 
     for field in fields(value):
         if field.name[0] == "_" and (
-                not hasattr(field.type, "__origin__")
-                or field.type.__origin__ is not ClassVar
+            not hasattr(field.type, "__origin__") or field.type.__origin__ is not ClassVar
         ):
             encoder.encode_string(to_java_field_name(field))
             context.serialize(getattr(value, field.name), None, encoder)
@@ -442,11 +405,11 @@ class DefaultValueSerializer(ValueSerializer):
         return self.serialize(*args, **kwargs)
 
     def serialize(
-            self,
-            value: Any,
-            type_name: Optional[str],
-            encoder: CBOREncoder,
-            context: SerializationContext,
+        self,
+        value: Any,
+        type_name: Optional[str],
+        encoder: CBOREncoder,
+        context: SerializationContext,
     ) -> None:
         if isinstance(value, (int, float, str, bool, decimal.Decimal)):
             encoder.encode(value)
@@ -455,11 +418,7 @@ class DefaultValueSerializer(ValueSerializer):
         elif isinstance(value, UUID):
             encoder.encode(value.bytes)
         elif isinstance(value, Enum):
-            if isinstance(value, JavaType.Primitive):
-                # FIXME implement type attribution support
-                encoder.encode(["org.openrewrite.java.tree.JavaType$Primitive", value.value])
-            else:
-                encoder.encode(value.value)
+            encoder.encode(value.value)
         elif isinstance(value, Path):
             encoder.encode(str(value))
         elif isinstance(value, (list, set, tuple)):
@@ -492,8 +451,8 @@ class DefaultValueSerializer(ValueSerializer):
                 encoder.encode_int(id)
                 for field in fields(value):
                     if field.name[0] == "_" and (
-                            not hasattr(field.type, "__origin__")
-                            or field.type.__origin__ is not ClassVar
+                        not hasattr(field.type, "__origin__")
+                        or field.type.__origin__ is not ClassVar
                     ):
                         encoder.encode_string(to_java_field_name(field))
                         context.serialize(getattr(value, field.name), None, encoder)
@@ -504,21 +463,19 @@ class DefaultValueSerializer(ValueSerializer):
         elif isinstance(value, complex):
             encoder.encode(str(value))
         else:
-            write_object_using_reflection(
-                value, type_name, False, encoder, context
-            )
+            write_object_using_reflection(value, type_name, False, encoder, context)
 
 
 def delegate_based_serializer(
-        delegate: Callable[[Any, Optional[str], CBOREncoder, SerializationContext], None]
+    delegate: Callable[[Any, Optional[str], CBOREncoder, SerializationContext], None],
 ) -> Type[ValueSerializer]:
     class DelegateBasedSerializer(ValueSerializer):
         def serialize(
-                self,
-                value: Any,
-                type_name: Optional[str],
-                encoder: CBOREncoder,
-                context: SerializationContext,
+            self,
+            value: Any,
+            type_name: Optional[str],
+            encoder: CBOREncoder,
+            context: SerializationContext,
         ) -> None:
             delegate(value, type_name, encoder, context)