openrewrite-remote 0.13.2__py3-none-any.whl → 0.13.3__py3-none-any.whl

rewrite_remote/handlers/recipe_install_handler.py

@@ -0,0 +1,152 @@
+import logging
+import socket
+from io import BytesIO
+from typing import List, TypedDict
+from cbor2 import dumps, CBORDecoder
+
+from rewrite_remote.remote_utils import COMMAND_END
+from rewrite_remote.remoting import OK, RemotingContext, RemotingMessageType
+from rewrite_remote.handlers.handler_helpers import respond_with_error
+from rewrite_remote.handlers.pypi_manager import PyPiManager, Source
+from rewrite_remote.handlers.types import PackageSource
+
+
+class RecipeInstallArgs(TypedDict):
+    package_id: str
+    package_version: str
+    include_default_repository: bool
+    package_sources: List[PackageSource]
+
+
+def decode_recipe_install_args(decoder: CBORDecoder) -> RecipeInstallArgs:
+    """
+    Decodes the arguments (order matters and must match the order encoded)
+    """
+    package_id = str(decoder.decode())
+    package_version = str(decoder.decode())
+    include_default_repository = bool(decoder.decode())
+    package_sources_data = decoder.decode()
+
+    if not isinstance(package_sources_data, list):
+        raise ValueError("package_sources_data is not a list")
+
+    package_sources = []
+    if package_sources_data and len(package_sources_data) > 0:
+        package_sources = [
+            PackageSource(source=ps.get("source"), credential=ps.get("credential"))
+            for ps in package_sources_data
+        ]
+
+    return {
+        "package_id": package_id,
+        "package_version": package_version,
+        "include_default_repository": include_default_repository,
+        "package_sources": package_sources,
+    }
+
+
+# Main command handler with the specified signature
+def recipe_install_handler(
+    stream: BytesIO, sock: socket.socket, remoting_ctx: RemotingContext
+) -> None:
+    remoting_ctx.reset()
+
+    # 1. Read input from stream
+    try:
+        data = stream.read()
+        decoder = CBORDecoder(BytesIO(data))
+        args = decode_recipe_install_args(decoder)
+        package_id = args.get("package_id")
+        package_version = args.get("package_version")
+        include_default_repository = args.get("include_default_repository")
+        package_sources = args.get("package_sources")
+    except Exception as e:  # pylint: disable=broad-except
+        respond_with_error(f"Failed to decode arguments: {e}", sock)
+        return
+
+    if package_id is None:
+        respond_with_error("package_id is required", sock)
+        return
+
+    if package_version is None:
+        respond_with_error("package_version is required", sock)
+        return
+
+    if package_sources is None:
+        respond_with_error("package_sources is required", sock)
+        return
+
+    if include_default_repository is None:
+        respond_with_error("include_default_repository is required", sock)
+        return
+
+    # 2. Log the request
+    logging.info(
+        f"""[Server] Handling install-recipe request: {{
+        packageId: {package_id},
+        packageVersion: {package_version},
+        packageSources: {package_sources},
+        includeDefaultRepository: {include_default_repository},
+    }}"""
+    )
+
+    # 3. Validate sources
+    sources: List[Source] = [
+        Source(
+            source=ps.source,
+            username=ps.credential.get("username") if ps.credential else None,
+            password=ps.credential.get("password") if ps.credential else None,
+            token=ps.credential.get("token") if ps.credential else None,
+        )
+        for ps in package_sources or []
+    ]
+
+    valid_source = PyPiManager.find_valid_source(
+        package_id, package_version, sources, include_default_repository
+    )
+
+    if not valid_source:
+        respond_with_error("No valid sources found", sock)
+        return
+
+    # 4. Install the recipe
+    try:
+        installable_recipes = PyPiManager.install_package(
+            package_id, package_version, valid_source
+        )
+    except Exception as e:  # pylint: disable=broad-except
+        respond_with_error(f"Failed to install package: {e}", sock)
+        return
+
+    # 5. Log the result
+    logging.info(
+        "[Server] Found %d recipe(s) for package %s",
+        len(installable_recipes.recipes),
+        package_id,
+    )
+    for recipe in installable_recipes.recipes:
+        logging.info("  Resolved recipe %s from %s", recipe, valid_source.source)
+
+    # 6. Write response to stream
+    response = {
+        "recipes": [
+            {
+                "name": recipe.name,
+                "source": recipe.source,
+                "options": [],
+            }
+            for recipe in installable_recipes.recipes
+        ],
+        "repository": installable_recipes.source,
+        "version": installable_recipes.version,
+    }
+
+    # Encode the response using CBOR
+    encoded_response = b""
+    encoded_response += dumps(RemotingMessageType.Response)
+    encoded_response += dumps(OK)
+    encoded_response += dumps(response)
+    encoded_response += COMMAND_END
+    sock.sendall(encoded_response)
+
+    logging.info("[Server] Request completed.")

rewrite_remote/handlers/run_recipe_load_and_visitor_handler.py

@@ -0,0 +1,135 @@
+import logging
+import socket
+from io import BytesIO
+from typing import List, TypedDict, Any
+from cbor2 import dumps, CBORDecoder
+
+from rewrite_remote.handlers.pypi_manager import Option, PyPiManager
+from rewrite_remote.remoting import (
+    OK,
+    RemotingExecutionContextView,
+    RemotingMessageType,
+    RemotingMessenger,
+)
+
+from rewrite_remote.remote_utils import COMMAND_END
+from rewrite_remote.remoting import OK, RemotingContext, RemotingMessageType
+from rewrite_remote.handlers.handler_helpers import respond_with_error
+
+from rewrite import InMemoryExecutionContext
+
+
+class RunRecipeLoadAndVisitorArgs(TypedDict):
+    recipe_name: str
+    recipe_source: str
+    recipe_options: Any  # List[recipeOption]
+
+
+def decode_run_recipe_load_and_visitor_args(
+    decoder: CBORDecoder,
+) -> RunRecipeLoadAndVisitorArgs:
+    """
+    Decodes the arguments (order matters and must match the order encoded)
+    """
+    recipe_name = str(decoder.decode())
+    recipe_source = str(decoder.decode())
+    recipe_options = decoder.decode()
+
+    return {
+        "recipe_name": recipe_name,
+        "recipe_source": recipe_source,
+        "recipe_options": recipe_options,
+    }
+
+
+def run_recipe_load_and_visitor_handler(
+    stream: BytesIO, sock: socket.socket, remoting_ctx: RemotingContext
+) -> None:
+    remoting_ctx.reset()
+
+    # Read input from stream
+    try:
+        data = stream.read()
+        decoder = CBORDecoder(BytesIO(data))
+        args = decode_run_recipe_load_and_visitor_args(decoder)
+        recipe_name = args.get("recipe_name")
+        recipe_source = args.get("recipe_source")
+        recipe_options: List[Option] = args.get("recipe_options") or []
+    except Exception as e:  # pylint: disable=broad-except
+        respond_with_error(f"Failed to decode arguments: {e}", sock)
+        return
+
+    if recipe_name is None:
+        respond_with_error("recipe_name is required", sock)
+        return
+
+    if recipe_source is None:
+        respond_with_error("recipe_source is required", sock)
+        return
+
+    # Log the request
+    logging.info(
+        f"""[Server] Handling run-recipe-load-and-visitor request: {{
+        recipe_name: {recipe_name},
+        recipe_source: {recipe_source},
+        recipe_options: {recipe_options},
+    }}"""
+    )
+
+    # Receive the tree
+    if not hasattr(RemotingMessenger, "_state"):
+        RemotingMessenger._state = None
+
+    if not hasattr(RemotingMessenger, "_context"):
+        RemotingMessenger._context = remoting_ctx
+
+    received = None
+    try:
+        received = RemotingMessenger.receive_tree(
+            RemotingMessenger, stream, RemotingMessenger._state
+        )
+    except Exception as e:
+        # logging.error(f"Failed to receive tree: {e}")
+        respond_with_error(f"Failed to receive tree: {e}", sock)
+        return
+
+    # Set the execution context
+    ctx = InMemoryExecutionContext()
+    RemotingExecutionContextView.view(ctx).remoting_context = remoting_ctx
+
+    try:
+        recipe_instance = PyPiManager.load_recipe(
+            recipe_name, recipe_source, recipe_options
+        )
+    except Exception as e:
+        respond_with_error(f"Failed to load recipe: {e}", sock)
+        return
+
+    # 4. Run the recipe
+    try:
+        tree_visitor = recipe_instance.get_visitor()
+
+        if not hasattr(tree_visitor, "visit") or not callable(tree_visitor.visit):
+            raise ValueError("Visitor does not have a visit method")
+
+        RemotingMessenger._state = tree_visitor.visit(received, ctx)
+
+        if RemotingMessenger._state is None:
+            raise ValueError("RemotingMessenger.state cannot be None")
+    except Exception as e:
+        respond_with_error(f"Failed to process input data with recipe: {e}", sock)
+        return
+
+    # 5. Write the response
+    response_encoder = BytesIO()
+    RemotingMessenger.send_tree(
+        remoting_ctx, response_encoder, RemotingMessenger._state, received
+    )
+
+    encoded_response = b""
+    encoded_response += dumps(RemotingMessageType.Response)
+    encoded_response += dumps(OK)
+    encoded_response += COMMAND_END
+    sock.sendall(encoded_response)
+
+    logging.info("[Server] Request completed.")

rewrite_remote/handlers/types.py

@@ -0,0 +1,30 @@
+from typing import List, Optional, Dict, Any
+from dataclasses import dataclass, field
+
+
+@dataclass
+class PackageSource:
+    source: str = ""
+    credential: Optional[Dict[str, str]] = None
+
+
+@dataclass
+class RecipeOption:
+    name: str
+    type: str
+    required: bool
+    value: Optional[Any]
+
+
+@dataclass
+class Recipe:
+    name: str
+    source: str
+    options: List[RecipeOption] = field(default_factory=list)
+
+
+@dataclass
+class RecipeInstallResponse:
+    recipes: List[Recipe]
+    repository: str
+    version: str

rewrite_remote/receiver.py

@@ -15,13 +15,13 @@ from typing import (
     cast,
     Iterable,
     Any,
-    Tuple,
-    get_args,
     TYPE_CHECKING,
+    get_args,
+    get_origin,
 )
 from uuid import UUID
 
-import cbor2
+from _cbor2 import break_marker
 from cbor2 import CBORDecoder
 from cbor2._decoder import major_decoders
 from rewrite import (
@@ -34,8 +34,8 @@ from rewrite import (
 )
 from rewrite import Tree, TreeVisitor, Cursor, FileAttributes
 
-from .event import DiffEvent, EventType
 from . import remote_utils, type_utils
+from .event import DiffEvent, EventType
 
 if TYPE_CHECKING:
     from .remoting import RemotingContext
@@ -198,17 +198,11 @@ class ReceiverContext:
         ReceiverContext.Registry[type_] = receiver_factory
 
 
-class ValueDeserializer(Protocol):
-    def deserialize(
-        self,
-        type_: Optional[Type[Any]],
-        reader: CBORDecoder,
-        context: "DeserializationContext",
-    ) -> Optional[Any]: ...
+ValueDeserializer = Callable[[Type[Any], CBORDecoder, "DeserializationContext"], Optional[Any]]
 
 
 class DefaultValueDeserializer(ValueDeserializer):
-    def deserialize(
+    def __call__(
         self,
         expected_type: Optional[Type[Any]],
         reader: CBORDecoder,
@@ -224,15 +218,15 @@ class DefaultValueDeserializer(ValueDeserializer):
 class DeserializationContext:
     DefaultDeserializer = DefaultValueDeserializer()
 
-    value_deserializers: List[Tuple[Type, ValueDeserializer]]
+    value_deserializers: Dict[str, ValueDeserializer]
 
     def __init__(
         self,
         remoting_context: "RemotingContext",
-        value_deserializers: Optional[List[Tuple[Type, ValueDeserializer]]] = None,
+        value_deserializers: Optional[Dict[Type, ValueDeserializer]] = None,
     ):
         self.remoting_context = remoting_context
-        self.value_deserializers = value_deserializers or []
+        self.value_deserializers = value_deserializers or {}
 
     def deserialize(self, expected_type: Type[Any], decoder: CBORDecoder) -> Any:
         if expected_type == UUID:
@@ -256,21 +250,36 @@ class DeserializationContext:
         if expected_type == float:
             return decoder.decode()
 
-        if issubclass(expected_type, Enum):
+        if isinstance(expected_type, type) and issubclass(expected_type, Enum):
             return expected_type(decoder.decode())
 
         initial_byte = decoder.read(1)[0]
         major_type = initial_byte >> 5
         subtype = initial_byte & 31
+        concrete_type = None
 
         # Object ID for Marker, JavaType, etc.
         if major_type == 0:
-            obj_id = major_decoders[major_type](decoder, subtype)
+            obj_id = decoder.decode_uint(subtype)
             return self.remoting_context.get_object(obj_id)
 
+        elif major_type == 1:
+            return decoder.decode_negint(subtype)
+
         # arrays
         elif major_type == 4:
-            pass
+            if get_origin(expected_type) in (List, list):
+                expected_elem_type = get_args(expected_type)[0]
+                array = []
+                if subtype != 31:
+                    for _ in range(subtype):
+                        array.append(self.deserialize(expected_elem_type, decoder))
+                else:
+                    while not (value := self.deserialize(expected_elem_type, decoder)) == break_marker:
+                        array.append(value)
+                return array
+            else:
+                concrete_type = decoder.decode()
 
         # objects
         elif major_type == 5:
@@ -278,6 +287,7 @@ class DeserializationContext:
             assert field_name == "@c"
             concrete_type = decoder.decode()
 
+        if concrete_type:
             if concrete_type == "org.openrewrite.marker.SearchResult":
                 field_name = decoder.decode()
                 if field_name == "description":
@@ -291,8 +301,12 @@ class DeserializationContext:
                     id_ = UUID(bytes=decoder.decode())
                 return SearchResult(id_, desc)
 
-            for type, value_deserializer in self.value_deserializers:
-                pass
+            if (deser := self.value_deserializers.get(concrete_type)):
+                return deser(concrete_type, decoder, self)
+
+            for type_, value_deserializer in self.value_deserializers.items():
+                if type_ == concrete_type:
+                    return value_deserializer(concrete_type, decoder, self)
 
             initial_byte = decoder.read(1)[0]
             major_type = initial_byte >> 5
@@ -305,6 +319,10 @@ class DeserializationContext:
                 subtype = initial_byte & 31
             pass
 
+        elif major_type == 3:
+            return decoder.decode_string(subtype)
+        elif major_type == 2:
+            return decoder.decode_bytestring(subtype)
         else:
             return major_decoders[major_type](decoder, subtype)
 
@@ -368,7 +386,7 @@ class DeserializationContext:
             decoder.read_array_start()
             concrete_type = decoder.read_text_string()
             actual_type = type_utils.get_type(concrete_type)
-            for type_, deserializer in self.value_deserializers:
+            for type_, deserializer in self.value_deserializers.items():
                 if issubclass(actual_type, type_):
                     return deserializer.deserialize(actual_type, decoder, self)
 
@@ -404,7 +422,7 @@ class DeserializationContext:
                 raise NotImplementedError("Expected @c key")
             concrete_type = decoder.read_text_string()
             actual_type = type_utils.get_type(concrete_type)
-            for type_, deserializer in self.value_deserializers:
+            for type_, deserializer in self.value_deserializers.items():
                 if issubclass(actual_type, type_):
                     return deserializer.deserialize(actual_type, decoder, self)
         raise NotImplementedError(
@@ -418,7 +436,7 @@ class JsonReceiver(TreeReceiver):
     def __init__(self, stream, context: DeserializationContext):
         super().__init__()
         self._stream = stream
-        self._decoder = cbor2.CBORDecoder(self._stream)
+        self._decoder = CBORDecoder(self._stream)
         self._context = context
         self._count = 0
 

rewrite_remote/remoting.py

@@ -56,7 +56,7 @@ class RemotingContext:
     _remoting_thread_local = threading.local()
     _recipe_factories: Dict[str, Callable[[str, Dict[str, Any]], Recipe]] = {}
     _value_serializers: Dict[Type, ValueSerializer] = {}
-    _value_deserializers: List[Tuple[Type, ValueDeserializer]] = []
+    _value_deserializers: Dict[str, ValueDeserializer] = {}
     _object_to_id_map: Dict[int, int] = {}
     _id_to_object_map: Dict[int, Any] = {}
 
@@ -126,19 +126,12 @@ class RemotingContext:
         return self._recipe_factories[recipe_id](recipe_options)
 
     @classmethod
-    def register_value_serializer(cls, serializer: ValueSerializer) -> None:
-        cls._value_serializers[serializer.__annotations__["value"]] = serializer
+    def register_value_serializer(cls, type_: Type, serializer: ValueSerializer) -> None:
+        cls._value_serializers[type_] = 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"]:
-                cls._value_deserializers[i] = (type_, deserializer)
-                return
-        cls._value_deserializers.append(
-            (deserializer.__annotations__["data"], deserializer)
-        )
+    def register_value_deserializer(cls, type_name: str, deserializer: ValueDeserializer) -> None:
+        cls._value_deserializers[type_name] = deserializer
 
 
 class RemotingExecutionContextView(DelegatingExecutionContext):