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

{rewrite/remote → rewrite_remote}/sender.py

@@ -1,100 +1,147 @@
+# type: ignore
 from __future__ import absolute_import
 
 import decimal
 from abc import ABC, abstractmethod
 from dataclasses import fields
 from pathlib import Path
-from typing import Dict, Type, Any, Callable, ClassVar, TypeVar, Generic, Protocol, List, TYPE_CHECKING
+from typing import (
+    Any,
+    Callable,
+    ClassVar,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Protocol,
+    TYPE_CHECKING,
+    Type,
+    TypeVar,
+)
 from uuid import UUID
-
 import cbor2
 from cbor2 import CBOREncoder
-from rewrite import Tree, Markers, Marker, ParseErrorVisitor, Cursor, Style
-from rewrite.remote.type_utils import to_java_type_name, to_java_field_name, to_java_type_name_from_value
+
 from rewrite.visitor import TreeVisitor
+from rewrite import Tree, Markers, Marker, ParseErrorVisitor, Cursor, Style
 
-from . import remote_utils
-from .event import *
-from .remote_utils import Operation
+from rewrite_remote.type_utils import (
+    to_java_type_name,
+    to_java_field_name,
+    to_java_type_name_from_value,
+)
+
+from rewrite_remote import remote_utils
+from rewrite_remote.event import *
+from rewrite_remote.remote_utils import Operation
 
 if TYPE_CHECKING:
     from .remoting import RemotingContext
 
-A = TypeVar('A')
-T = TypeVar('T', bound=Tree)
-V = TypeVar('V')
-I = TypeVar('I')
+A = TypeVar("A")
+T = TypeVar("T", bound=Tree)
+V = TypeVar("V")
+I = TypeVar("I")
 
 
 class Sender(Protocol):
-    def send(self, after: T, before: Optional[T], ctx: 'SenderContext') -> None:
-        ...
+    def send(self, after: T, before: Optional[T], ctx: "SenderContext") -> None: ...
 
 
 class OmniSender(Sender):
-    def send(self, after: Tree, before: Optional[Tree], ctx: 'SenderContext'):
+    def send(self, after: Tree, before: Optional[Tree], ctx: "SenderContext") -> None:
         sender = ctx.new_sender(type(after))
         sender.send(after, before, ctx)
 
 
 class TreeSender(Protocol):
-    def send_node(self, diff_event: DiffEvent, visitor: Callable[['TreeSender'], None]):
-        ...
+    def send_node(
+        self, diff_event: DiffEvent, visitor: Callable[["TreeSender"], None]
+    ) -> None: ...
 
-    def send_value(self, diff_event: DiffEvent):
-        ...
+    def send_value(self, diff_event: DiffEvent) -> None: ...
 
-    def flush(self):
-        ...
+    def flush(self) -> None: ...
 
 
 class SenderContext(Generic[T]):
-    Registry: Dict[Type, Callable[[], Sender]] = {}
-
-    def __init__(self, sender: 'TreeSender', visitor: TreeVisitor = None, before: Optional[Any] = None):
+    """
+    Context for serialization of tree nodes and values.
+    """
+
+    Registry: Dict[Type[Any], Callable[[], Sender]] = {}
+
+    def __init__(
+        self,
+        sender: "TreeSender",
+        visitor: TreeVisitor = None,
+        before: Optional[Any] = None,
+    ):
         self.sender = sender
         self.visitor = visitor
         self.before = before
 
-    def new_sender(self, type_: Type) -> Sender:
+    def new_sender(self, type_: Type[Any]) -> Sender:
         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':
+    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):
+    def visit(
+        self,
+        consumer: Callable[[V, "SenderContext"], None],
+        after: V,
+        before: Optional[V] = None,
+    ) -> None:
         save_before = self.before
         self.before = before
         consumer(after, self)
         self.before = save_before
 
-    def send_tree(self, after: T, ctx: 'SenderContext'):
+    def send_tree(self, after: T, ctx: "SenderContext[T]") -> None:
         after.accept(self.visitor, ctx)
 
-    def send_any_tree(self, after: T, before: Optional[T]):
+    def send_any_tree(self, after: T, before: Optional[T]) -> None:
         OmniSender().send(after, before, self)
 
-    def send_node(self, owner: A, extractor: Callable[[A], Optional[V]], details: Callable[[V, 'SenderContext'], None]):
-        self.send_node_internal(extractor(owner), extractor(self.before) if self.before is not None else None, details)
-
-    def send_value(self, owner, value_extractor: Callable[[T], Optional[V]]):
+    def send_node(
+        self,
+        owner: A,
+        extractor: Callable[[A], Optional[V]],
+        details: Callable[[V, "SenderContext[T]"], None],
+    ) -> None:
+        self.send_node_internal(
+            extractor(owner),
+            extractor(self.before) if self.before is not None else None,
+            details,
+        )
+
+    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)
 
-    def send_typed_value(self, owner: A, value_extractor: Callable[[A], V]):
+    def send_typed_value(self, owner: A, value_extractor: Callable[[A], V]) -> None:
         after_value = value_extractor(owner)
         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:
@@ -105,12 +152,21 @@ 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
-
-    def send_nodes(self, owner: A, element_extractor: Callable[[A], List[V]],
-                   details: Callable[[V, 'SenderContext'], None], id_function: Callable[[V], I]):
+        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],
+    ) -> 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:
@@ -122,17 +178,26 @@ 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))
-                }[op]()
+                    Operation.Move: lambda: NotImplementedError(
+                        "Unexpected operation: " + str(op)
+                    ),
+                }[op](),
             )
 
             if before_list is not None:
                 self.sender.send_value(DiffEvent(EventType.EndList, None, None))
 
-    def send_values(self, owner, value_extractor: Callable[[T], List[V]], id_function: Callable[[V], I]):
+    def send_values(
+        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
 
@@ -149,14 +214,21 @@ 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))
-                }[op]()
+                    Operation.Move: lambda: NotImplementedError(
+                        "Unexpected operation: " + str(op)
+                    ),
+                }[op](),
             )
 
             if before_list is not None:
                 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]):
+    def send_typed_values(
+        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
 
@@ -173,22 +245,24 @@ 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))
-                }[op]()
+                    Operation.Move: lambda: NotImplementedError(
+                        "Unexpected operation: " + str(op)
+                    ),
+                }[op](),
             )
 
             if before_list is not None:
                 self.sender.send_value(DiffEvent(EventType.EndList, None, None))
 
-    def send_markers(self, markers: Markers, ignore):
+    def send_markers(self, markers: Markers, ignore: Optional[bool]) -> None:
         self.send_value(markers, lambda ms: ms.id)
         self.send_values(markers, lambda ms: ms.markers, lambda ms: ms.id)
 
-    def send_tree_visitor(self, after: T, ctx: 'SenderContext'):
+    def send_tree_visitor(self, after: T, ctx: "SenderContext") -> None:
         after.accept(self.visitor, ctx)
 
     @staticmethod
-    def register(type_: Type, sender_factory: Callable[[], Sender]):
+    def register(type_: Type, sender_factory: Callable[[], Sender]) -> None:
         SenderContext.Registry[type_] = sender_factory
 
     @staticmethod
@@ -198,12 +272,18 @@ 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]):
+    def send_node_internal(
+        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)
@@ -212,11 +292,13 @@ class SenderContext(Generic[T]):
 
         self.sender.send_node(evt, lambda _: self.visit(details, after, before))
 
-    def send_value_internal(self, after: V, before: Optional[V]):
+    def send_value_internal(self, after: V, before: Optional[V]) -> None:
         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)
@@ -225,11 +307,13 @@ class SenderContext(Generic[T]):
 
         self.sender.send_value(evt)
 
-    def send_typed_value_internal(self, after: V, before: Optional[V]):
+    def send_typed_value_internal(self, after: V, before: Optional[V]) -> None:
         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)
@@ -240,12 +324,17 @@ class SenderContext(Generic[T]):
 
 
 class SerializationContext:
-    def __init__(self, remoting_context: 'RemotingContext',
-                 value_serializers: Optional[Dict[Type, 'ValueSerializer']] = None):
+    def __init__(
+        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):
+    def serialize(
+        self, value: Any, type_name: Optional[str], encoder: CBOREncoder
+    ) -> None:
         if value is None:
             encoder.encode_none(None)
             return
@@ -266,12 +355,14 @@ class JsonSender(TreeSender):
         self._context = context
         self._encoder = cbor2.CBOREncoder(self._stream)
 
-    def send_node(self, diff_event, visitor):
+    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] if diff_event.concrete_type is None else [diff_event.event_type.value,
-                                                                                        diff_event.concrete_type])
+                [diff_event.event_type.value]
+                if diff_event.concrete_type is None
+                else [diff_event.event_type.value, diff_event.concrete_type]
+            )
 
             if self.Debug:
                 print(f"SEND: {diff_event}")
@@ -287,12 +378,21 @@ 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)
-        elif diff_event.event_type in (EventType.Delete, EventType.NoChange, EventType.StartList, EventType.EndList):
+            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,
+        ):
             self._encoder.encode([diff_event.event_type.value])
         else:
             raise NotImplementedError()
@@ -304,23 +404,34 @@ class JsonSender(TreeSender):
         self._stream.flush()
 
 
-INDEFINITE_ARRAY_START = b'\x9f'
-INDEFINITE_MAP_START = b'\xbf'
-BREAK_MARKER = b'\xff'
+INDEFINITE_ARRAY_START = b"\x9f"
+INDEFINITE_MAP_START = b"\xbf"
+BREAK_MARKER = b"\xff"
 
 
 class ValueSerializer(ABC):
 
     @abstractmethod
-    def serialize(self, value: Any, type_name: Optional[str], writer: CBOREncoder, context: SerializationContext):
+    def serialize(
+        self,
+        value: Any,
+        type_name: Optional[str],
+        writer: CBOREncoder,
+        context: SerializationContext,
+    ) -> None:
         pass
 
     @staticmethod
-    def write_object_using_reflection(value: Any, type_name: Optional[str], with_id: bool,
-                                      encoder: CBOREncoder, context: SerializationContext):
-        if type(value).__qualname__ == 'JavaType.Primitive':
+    def write_object_using_reflection(
+        value: Any,
+        type_name: Optional[str],
+        with_id: bool,
+        encoder: CBOREncoder,
+        context: SerializationContext,
+    ) -> None:
+        if type(value).__qualname__ == "JavaType.Primitive":
             # FIXME implement type attribution support
-            encoder.encode(['org.openrewrite.java.tree.JavaType$Primitive', 0])
+            encoder.encode(["org.openrewrite.java.tree.JavaType$Primitive", 0])
             return
         if with_id and (id := context.remoting_context.try_get_id(value)):
             encoder.encode_int(id)
@@ -335,14 +446,23 @@ class ValueSerializer(ABC):
             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):
+            if field.name[0] == "_" and (
+                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)
         encoder.write(BREAK_MARKER)
 
 
 class DefaultValueSerializer(ValueSerializer):
-    def serialize(self, value: Any, type_name: Optional[str], encoder: CBOREncoder, context: SerializationContext):
+    def serialize(
+        self,
+        value: Any,
+        type_name: Optional[str],
+        encoder: CBOREncoder,
+        context: SerializationContext,
+    ) -> None:
         if isinstance(value, (int, float, str, bool, decimal.Decimal)):
             encoder.encode(value)
         elif value is None:
@@ -358,32 +478,34 @@ class DefaultValueSerializer(ValueSerializer):
             for item in value:
                 context.serialize(item, None, encoder)
         elif isinstance(value, Markers):
-            if (id := context.remoting_context.try_get_id(value)):
+            if id := context.remoting_context.try_get_id(value):
                 encoder.encode_int(id)
             else:
                 id = context.remoting_context.add(value)
                 encoder.encode_length(5, 3)
-                encoder.encode_string('@ref')
+                encoder.encode_string("@ref")
                 encoder.encode_int(id)
-                encoder.encode_string('id')
+                encoder.encode_string("id")
                 encoder.encode_uuid(value.id)
-                encoder.encode_string('markers')
+                encoder.encode_string("markers")
                 encoder.encode_length(4, len(value.markers))
                 for marker in value.markers:
                     context.serialize(marker, None, encoder)
         elif isinstance(value, (Marker, Style)):
-            if (id := context.remoting_context.try_get_id(value)):
+            if id := context.remoting_context.try_get_id(value):
                 encoder.encode_int(id)
             else:
                 id = context.remoting_context.add(value)
                 encoder.write(INDEFINITE_MAP_START)
                 encoder.encode_string("@c")
                 encoder.encode_string(to_java_type_name(type(value)))
-                encoder.encode_string('@ref')
+                encoder.encode_string("@ref")
                 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):
+                    if field.name[0] == "_" and (
+                        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)
                 encoder.write(BREAK_MARKER)
@@ -393,19 +515,29 @@ class DefaultValueSerializer(ValueSerializer):
         elif isinstance(value, complex):
             encoder.encode(str(value))
         else:
-            ValueSerializer.write_object_using_reflection(value, type_name, False, encoder, context)
+            ValueSerializer.write_object_using_reflection(
+                value, type_name, False, encoder, context
+            )
 
 
-def delegate_based_serializer(delegate: Callable[[Any, Optional[str], CBOREncoder, SerializationContext], None]):
+def delegate_based_serializer(
+    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):
+        def serialize(
+            self,
+            value: Any,
+            type_name: Optional[str],
+            encoder: CBOREncoder,
+            context: SerializationContext,
+        ) -> None:
             delegate(value, type_name, encoder, context)
 
     return DelegateBasedSerializer
 
 
 class ParseErrorSender(Sender):
-    def send(self, after, before, ctx):
+    def send(self, after: Tree, before: Optional[Tree], ctx: "SenderContext") -> None:
         visitor = self.Visitor()
         visitor.visit(after, ctx.fork(visitor, before))