openrewrite-remote 0.12.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

{rewrite/remote → rewrite_remote}/remote_utils.py

@@ -1,18 +1,30 @@
+# type: ignore
 import socket
 import struct
 from enum import Enum, auto
 from io import BytesIO
-from typing import Callable, List, Optional, TypeVar, Dict, TYPE_CHECKING, Iterable, cast, Type
+from typing import (
+    Callable,
+    List,
+    Optional,
+    TypeVar,
+    Dict,
+    TYPE_CHECKING,
+    Iterable,
+    cast,
+    Type,
+    Any,
+)
 
 from cbor2 import CBORDecoder, CBORDecodeValueError
 
 from .event import EventType
 
 if TYPE_CHECKING:
-    from .receiver import ReceiverContext, DetailsReceiver
+    from rewrite_remote.receiver import ReceiverContext, DetailsReceiver
 
-T = TypeVar('T')
-I = TypeVar('I')
+T = TypeVar("T")
+I = TypeVar("I")
 
 
 class Operation(Enum):
@@ -23,9 +35,11 @@ class Operation(Enum):
     Move = auto()
 
 
-def receive_nodes(before: Optional[List[T]],
-                  details: Callable[[Optional[T], Optional[str], 'ReceiverContext'], T],
-                  ctx: 'ReceiverContext') -> Optional[List[T]]:
+def receive_nodes(
+    before: Optional[List[T]],
+    details: Callable[[Optional[T], Optional[str], "ReceiverContext"], T],
+    ctx: "ReceiverContext",
+) -> Optional[List[T]]:
     list_event = ctx.receiver.receive_value(list)
     if list_event.event_type == EventType.NoChange:
         return before
@@ -33,26 +47,32 @@ def receive_nodes(before: Optional[List[T]],
         return None
     elif list_event.event_type == EventType.Add:
         after_size = list_event.msg
+        if after_size is None:
+            after_size = 0
         after = [None] * after_size  # Preallocate list
         for i in range(after_size):
             diff_event = ctx.receiver.receive_node()
             if diff_event.event_type == EventType.Add:
-                after[i] = details.receive_details(None, diff_event.concrete_type, ctx)
+                after[i] = details(None, diff_event.concrete_type, ctx)  # type: ignore
             elif diff_event.event_type == EventType.NoChange:
                 after[i] = None  # Or some default value
             else:
-                raise NotImplementedError(f"Unexpected operation: {diff_event.event_type}")
-        return after
+                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)
+        return _receive_updated_nodes(before, list_event.msg, details, ctx)  # type: ignore
     else:
         raise NotImplementedError(f"Unexpected operation: {list_event.event_type}")
 
 
-def _receive_updated_nodes(before: List[T],
-                           after_size: int,
-                           details: 'DetailsReceiver[T]',
-                           ctx: 'ReceiverContext') -> List[T]:
+def _receive_updated_nodes(
+    before: List[T],
+    after_size: int,
+    details: "DetailsReceiver[T]",
+    ctx: "ReceiverContext",
+) -> List[T]:
     modified = False
     after_list = before
     evt = ctx.receiver.receive_node()
@@ -77,7 +97,9 @@ def _receive_updated_nodes(before: List[T],
         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))
@@ -88,7 +110,9 @@ def _receive_updated_nodes(before: List[T],
     return after_list[:after_size] if len(after_list) > after_size else after_list
 
 
-def receive_values(before: Optional[List[T]], type: Type, ctx: 'ReceiverContext') -> Optional[List[T]]:
+def receive_values(
+    before: Optional[List[T]], type: Type[Any], ctx: "ReceiverContext"
+) -> Optional[List[T]]:
     list_event = ctx.receiver.receive_value(list)
     if list_event.event_type == EventType.NoChange:
         return before
@@ -96,6 +120,8 @@ def receive_values(before: Optional[List[T]], type: Type, ctx: 'ReceiverContext'
         return None
     elif list_event.event_type == EventType.Add:
         after_size = list_event.msg
+        if after_size is None:
+            after_size = 0
         after = [None] * after_size  # Preallocate list
         for i in range(after_size):
             diff_event = ctx.receiver.receive_value(type)
@@ -104,15 +130,19 @@ def receive_values(before: Optional[List[T]], type: Type, ctx: 'ReceiverContext'
             elif diff_event.event_type == EventType.NoChange:
                 after[i] = None  # Or some default value
             else:
-                raise NotImplementedError(f"Unexpected operation: {diff_event.event_type}")
-        return after
+                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)
+        return _receive_updated_values(before, list_event.msg, type, ctx)  # type: ignore
     else:
         raise NotImplementedError(f"Unexpected operation: {list_event.event_type}")
 
 
-def _receive_updated_values(before: List[T], after_size: int, type: Type, ctx: 'ReceiverContext') -> List[T]:
+def _receive_updated_values(
+    before: List[T], after_size: int, type: Type[Any], ctx: "ReceiverContext"
+) -> List[T]:
     modified = False
     after_list = before
     evt = ctx.receiver.receive_node()
@@ -137,7 +167,7 @@ def _receive_updated_values(before: List[T], after_size: int, type: Type, ctx: '
         elif evt.event_type == EventType.Delete:
             before_idx += 1
         elif evt.event_type in (EventType.Update, EventType.Add):
-            after_list.append(evt.msg)
+            after_list.append(cast(T, evt.msg))
             if evt.event_type == EventType.Update:
                 before_idx += 1
 
@@ -150,17 +180,19 @@ def _receive_updated_values(before: List[T], after_size: int, type: Type, ctx: '
 def _copy_range(before: Iterable[T], j: int) -> List[T]:
     if isinstance(before, list):
         return before[:j]
-    elif hasattr(before, 'getrange'):  # If the object has a 'getrange' method (e.g., an immutable list)
-        return before.getrange(0, j)
+    elif hasattr(
+        before, "getrange"
+    ):  # If the object has a 'getrange' method (e.g., an immutable list)
+        return cast(List[T], before.getrange(0, j))
     else:
         return list(before)[:j]
 
 
 def calculate_list_diff(
-        before: List[T],
-        after: List[T],
-        id_function: Callable[[T], I],
-        consumer: Callable[[Operation, int, int, Optional[T], Optional[T]], None]
+    before: List[T],
+    after: List[T],
+    id_function: Callable[[T], I],
+    consumer: Callable[[Operation, int, int, Optional[T], Optional[T]], None],
 ) -> None:
     before_idx, after_idx = 0, 0
     before_size, after_size = len(before), len(after)
@@ -183,7 +215,7 @@ def calculate_list_diff(
                 before_idx,
                 after_idx,
                 before[before_idx],
-                after[after_idx]
+                after[after_idx],
             )
             before_idx += 1
             after_idx += 1
@@ -197,7 +229,7 @@ def calculate_list_diff(
                     before_idx,
                     after_idx,
                     before[before_idx],
-                    after[after_idx]
+                    after[after_idx],
                 )
                 before_idx += 1
                 after_idx += 1
@@ -215,9 +247,7 @@ def calculate_list_diff(
 
 
 def create_index_map(
-        lst: List[T],
-        from_index: int,
-        id_function: Callable[[T], I]
+    lst: List[T], from_index: int, id_function: Callable[[T], I]
 ) -> Dict[I, int]:
     result = {}
     for i in range(from_index, len(lst)):
@@ -225,7 +255,9 @@ def create_index_map(
     return result
 
 
-def _decode_length(decoder: CBORDecoder, subtype: int, allow_indefinite: bool = False):
+def _decode_length(
+    decoder: CBORDecoder, subtype: int, allow_indefinite: bool = False
+) -> Optional[int]:
     if subtype < 24:
         return subtype
     elif subtype == 24:
@@ -241,7 +273,8 @@ def _decode_length(decoder: CBORDecoder, subtype: int, allow_indefinite: bool =
     else:
         raise CBORDecodeValueError(f"unknown unsigned integer subtype 0x{subtype:x}")
 
-def decode_array_start(decoder: CBORDecoder):
+
+def decode_array_start(decoder: CBORDecoder) -> Optional[int]:
     initial_byte = decoder.read(1)[0]
     major_type = initial_byte >> 5
     assert major_type == 4
@@ -253,6 +286,7 @@ COPY_BUFFER_SIZE = 4096
 COPY_BUFFER = bytearray(COPY_BUFFER_SIZE)
 COMMAND_END = bytes([0x81, 0x17])
 
+
 def read_to_command_end(sock: socket.socket) -> BytesIO:
     memory_stream = BytesIO()
 
@@ -262,15 +296,19 @@ def read_to_command_end(sock: socket.socket) -> BytesIO:
             if bytes_read == 0:
                 break
             memory_stream.write(COPY_BUFFER[:bytes_read])
-            if (bytes_read > 1 and
-                    COPY_BUFFER[bytes_read - 2] == COMMAND_END[0] and
-                    COPY_BUFFER[bytes_read - 1] == COMMAND_END[1]):
+            if (
+                bytes_read > 1
+                and COPY_BUFFER[bytes_read - 2] == COMMAND_END[0]
+                and COPY_BUFFER[bytes_read - 1] == COMMAND_END[1]
+            ):
                 break
             elif bytes_read == 1:
                 original_position = memory_stream.tell()
                 memory_stream.seek(-2, 1)  # Move back by 2 bytes
-                if (memory_stream.read(1)[0] == COMMAND_END[0] and
-                        memory_stream.read(1)[0] == COMMAND_END[1]):
+                if (
+                    memory_stream.read(1)[0] == COMMAND_END[0]
+                    and memory_stream.read(1)[0] == COMMAND_END[1]
+                ):
                     memory_stream.seek(original_position)
                     break
     except socket.error as e:

{rewrite/remote → rewrite_remote}/remoting.py

@@ -1,3 +1,4 @@
+# type: ignore
 from __future__ import absolute_import
 
 import socket
@@ -6,19 +7,52 @@ import traceback
 from dataclasses import dataclass
 from io import BytesIO
 from threading import Lock
-from typing import Any, Dict, Optional, Type, List, Tuple, Callable, cast, Union, BinaryIO
+from typing import (
+    Any,
+    Dict,
+    Optional,
+    Type,
+    List,
+    Tuple,
+    Callable,
+    cast,
+    Union,
+    BinaryIO,
+)
 
 import cbor2
 from cbor2 import dumps, loads, load
-from rewrite import Recipe, InMemoryExecutionContext, Cursor, PrinterFactory, TreeVisitor, Tree
+
+from rewrite import (
+    Recipe,
+    InMemoryExecutionContext,
+    Cursor,
+    PrinterFactory,
+    TreeVisitor,
+    Tree,
+)
 from rewrite.execution import DelegatingExecutionContext
-from rewrite.remote import ValueSerializer, ValueDeserializer, SenderContext, ReceiverContext, JsonSender, JsonReceiver, \
-    SerializationContext, DeserializationContext, remote_utils
 from rewrite.tree import PrintOutputCapture, P
 from rewrite.visitor import T
 
+from rewrite_remote import (
+    ValueSerializer,
+    ValueDeserializer,
+    SenderContext,
+    ReceiverContext,
+    JsonSender,
+    JsonReceiver,
+    SerializationContext,
+    DeserializationContext,
+    remote_utils,
+)
+
 
 class RemotingContext:
+    """
+    A context for remoting operations.
+    """
+
     _remoting_thread_local = threading.local()
     _recipe_factories: Dict[str, Callable[[str, Dict[str, Any]], Recipe]] = {}
     _value_serializers: Dict[Type, ValueSerializer] = {}
@@ -30,8 +64,8 @@ class RemotingContext:
         self._client = None
 
     @classmethod
-    def current(cls) -> 'RemotingContext':
-        result = getattr(cls._remoting_thread_local, 'context', None)
+    def current(cls) -> "RemotingContext":
+        result = getattr(cls._remoting_thread_local, "context", None)
         if result is None:
             raise ValueError("No RemotingContext has been set")
         return result
@@ -41,10 +75,10 @@ class RemotingContext:
         cls._remoting_thread_local.context = self
 
     @property
-    def client(self) -> Optional['RemotingClient']:
+    def client(self) -> Optional["RemotingClient"]:
         return self._client
 
-    def connect(self, sock: Any) -> 'RemotingContext':
+    def connect(self, sock: Any) -> "RemotingContext":
         self._client = RemotingClient(self, sock)
         return self
 
@@ -71,30 +105,40 @@ class RemotingContext:
         self._object_to_id_map.clear()
         self._id_to_object_map.clear()
 
-    def new_sender_context(self, output_stream: Any) -> 'SenderContext':
-        return SenderContext(JsonSender(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)))
-
-    def copy(self) -> 'RemotingContext':
+    def new_sender_context(self, output_stream: Any) -> "SenderContext":
+        return SenderContext(
+            JsonSender(
+                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)
+            )
+        )
+
+    def copy(self) -> "RemotingContext":
         return RemotingContext()
 
-    def new_recipe(self, recipe_id: str, recipe_options: Any) -> 'Recipe':
+    def new_recipe(self, recipe_id: str, recipe_options: Any) -> "Recipe":
         return self._recipe_factories[recipe_id](recipe_options)
 
     @classmethod
     def register_value_serializer(cls, serializer: ValueSerializer) -> None:
-        cls._value_serializers[serializer.__annotations__['value']] = serializer
+        cls._value_serializers[serializer.__annotations__["value"]] = serializer
 
     @classmethod
     def register_value_deserializer(cls, deserializer: ValueDeserializer) -> None:
         for i in range(len(cls._value_deserializers)):
             type_ = cls._value_deserializers[i][0]
-            if type_ == deserializer.__annotations__['data']:
+            if type_ == deserializer.__annotations__["data"]:
                 cls._value_deserializers[i] = (type_, deserializer)
                 return
-        cls._value_deserializers.append((deserializer.__annotations__['data'], deserializer))
+        cls._value_deserializers.append(
+            (deserializer.__annotations__["data"], deserializer)
+        )
 
 
 class RemotingExecutionContextView(DelegatingExecutionContext):
@@ -110,12 +154,12 @@ class RemotingExecutionContextView(DelegatingExecutionContext):
 
     @property
     def remoting_context(self) -> RemotingContext:
-        return self._delegate.get_message('remoting', RemotingContext.current())
+        return self._delegate.get_message("remoting", RemotingContext.current())
 
     @remoting_context.setter
     def remoting_context(self, value: RemotingContext):
         value.set_current()
-        self._delegate.put_message('remoting', value)
+        self._delegate.put_message("remoting", value)
 
 
 OK = 0
@@ -124,8 +168,13 @@ ERROR = 1
 
 class RemotingMessenger:
 
-    def __init__(self, context: RemotingContext,
-                 additional_handlers: Dict[str, Callable[[BytesIO, socket.socket, RemotingContext], Any]] = None):
+    def __init__(
+        self,
+        context: RemotingContext,
+        additional_handlers: Dict[
+            str, Callable[[BytesIO, socket.socket, RemotingContext], Any]
+        ] = None,
+    ):
         self._context = context
         self._additional_handlers = additional_handlers or {}
         self._recipes = []
@@ -190,7 +239,7 @@ class RemotingMessenger:
     def handle_run_recipe_visitor_command(self, stream: BytesIO, sock: socket.socket):
         recipe_index = cbor2.load(stream)
         recipe = self._recipes[recipe_index]
-        received = self.receive_tree(sock, self._state)
+        received = self.receive_tree(stream, self._state)
 
         ctx = InMemoryExecutionContext()
         RemotingExecutionContextView.view(ctx).remoting_context = self._context
@@ -207,7 +256,9 @@ 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)
@@ -232,11 +283,12 @@ class RemotingMessenger:
 
     @staticmethod
     def send_end_message(sock):
-        sock.sendall(b'\x81\x17')
+        sock.sendall(b"\x81\x17")
 
     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)))
+        self.__send_request_stream(
+            sock, "print", lambda s: self.send_tree(s, cast(Tree, cursor.value))
+        )
         if self.recv_byte(sock) != RemotingMessageType.Response:
             raise ValueError("Unexpected message type.")
         if self.recv_byte(sock) != 0:
@@ -246,22 +298,35 @@ class RemotingMessenger:
         # end = load(data)  # end
         return print_msg
 
-    def send_tree(self, sock: Union[BinaryIO, socket.socket], after: Tree, before: Optional[Tree] = None):
+    def send_tree(
+        self,
+        sock: 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()))
 
-    def receive_tree(self, data: Union[BinaryIO, socket.socket], before: Optional[Tree] = None):
-        receiver_context = self._context.new_receiver_context(cbor2.load(data))
+    def receive_tree(
+        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)
 
-    def send_run_recipe_request(self, sock: socket.socket, recipe, options: dict, source_files: list):
-        self.__send_request_stream(sock, "run-recipe", lambda s: (
-            sock.sendall(dumps(recipe)),
-            sock.sendall(dumps(options)),
-            sock.sendall(dumps(len(source_files))),
-            *[self.send_tree(sock, sf, None) for sf in source_files]
-        ))
+    def send_run_recipe_request(
+        self, sock: socket.socket, recipe, options: dict, source_files: list
+    ):
+        self.__send_request_stream(
+            sock,
+            "run-recipe",
+            lambda s: (
+                sock.sendall(dumps(recipe)),
+                sock.sendall(dumps(options)),
+                sock.sendall(dumps(len(source_files))),
+                *[self.send_tree(sock, sf, None) for sf in source_files],
+            ),
+        )
         while self.recv_byte(sock) == RemotingMessageType.Request:
             self.process_request(sock)
         if self.recv_byte(sock) != RemotingMessageType.Response:
@@ -285,7 +350,7 @@ class RemotingMessenger:
         return sock.recv(1)[0]
 
     def recv_all(self, sock, buffer_size=4096):
-        data = b''
+        data = b""
         while True:
             part = sock.recv(buffer_size)
             data += part
@@ -308,28 +373,35 @@ class RemotingClient:
     def close(self) -> None:
         self._socket.close()
 
-    def hello(self):
+    def hello(self) -> None:
         with self._lock:
             self._messenger.send_request(self._socket, "hello")
 
-    def print(self, cursor):
+    def print(self, cursor: Cursor) -> Any:
         with self._lock:
             return self._messenger.send_print_request(self._socket, cursor)
 
-    def reset(self):
+    def reset(self) -> None:
         with self._lock:
             self._messenger.send_reset_request(self._socket)
 
     def run_recipe(self, recipe, options: dict, source_files: list):
         with self._lock:
-            return self._messenger.send_run_recipe_request(self._socket, recipe, options, source_files)
+            return self._messenger.send_run_recipe_request(
+                self._socket, recipe, options, source_files
+            )
 
 
 @dataclass
 class RemotePrinter(TreeVisitor[Any, PrintOutputCapture[P]]):
     _client: RemotingClient
 
-    def visit(self, tree: Optional[Tree], p: PrintOutputCapture[P], parent: Optional[Cursor] = None) -> Optional[T]:
+    def visit(
+        self,
+        tree: Optional[Tree],
+        p: PrintOutputCapture[P],
+        parent: Optional[Cursor] = None,
+    ) -> Optional[T]:
         self.cursor = Cursor(parent, tree)
         p.append(self._client.print(self.cursor))
         self.cursor = self.cursor.parent