add register_object

Author: Mahad-10

tests/codec_test.py

@@ -100,3 +100,82 @@ def event_handler(event: Event):
     session.publish_object("io.xconn.object", String("hello"))
 
     session.leave()
+
+
+def test_register_object_one_param_with_return_type():
+    session = connect_anonymous("ws://localhost:8080/ws", "realm1")
+    session.set_payload_codec(ProtobufCodec())
+
+    def create_profile_handler(prof: ProfileCreate) -> ProfileGet:
+        return ProfileGet(
+            id="356",
+            username=prof.username,
+            email=prof.email,
+            age=prof.age,
+            created_at="2025-10-30T17:00:00Z",
+        )
+
+    session.register_object("io.xconn.profile.create", create_profile_handler)
+
+    profile_create = ProfileCreate(username="john", email="[email protected]", age=25)
+    result = session.call("io.xconn.profile.create", [profile_create.SerializeToString()])
+
+    profile = ProfileGet()
+    profile.ParseFromString(result.args[0])
+
+    assert profile.id == "356"
+    assert profile.username == "john"
+    assert profile.email == "[email protected]"
+    assert profile.age == 25
+    assert profile.created_at == "2025-10-30T17:00:00Z"
+
+    session.leave()
+
+
+def test_register_object_no_param():
+    session = connect_anonymous("ws://localhost:8080/ws", "realm1")
+    session.set_payload_codec(ProtobufCodec())
+
+    options = {"flag": False}
+
+    def invocation_handler() -> None:
+        options["flag"] = True
+
+    session.register_object("io.xconn.param.none", invocation_handler)
+
+    result = session.call("io.xconn.param.none")
+
+    assert options["flag"] is True
+    assert result.args is None
+    assert result.kwargs is None
+
+    session.leave()
+
+
+def test_register_object_no_param_with_return():
+    session = connect_anonymous("ws://localhost:8080/ws", "realm1")
+    session.set_payload_codec(ProtobufCodec())
+
+    def get_profile_handler() -> ProfileGet:
+        return ProfileGet(
+            id="636",
+            username="admin",
+            email="[email protected]",
+            age=30,
+            created_at="2025-10-30T17:00:00Z",
+        )
+
+    session.register_object("io.xconn.profile.get", get_profile_handler)
+
+    result = session.call("io.xconn.profile.get")
+
+    profile = ProfileGet()
+    profile.ParseFromString(result.args[0])
+
+    assert profile.id == "636"
+    assert profile.username == "admin"
+    assert profile.email == "[email protected]"
+    assert profile.age == 30
+    assert profile.created_at == "2025-10-30T17:00:00Z"
+
+    session.leave()

xconn/session.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import inspect
 from concurrent.futures import Future
 from threading import Thread
 from typing import Callable, Any, TypeVar, Type
@@ -231,6 +232,48 @@ def publish_object(self, topic: str, request: TReq):
         encoded = self._payload_codec.encode(request)
         return self.publish(topic, [encoded])
 
+    def register_object(
+        self,
+        procedure: str,
+        invocation_handler: Callable[[TReq], TRes | None] | Callable[[], TRes | None],
+    ):
+        if self._payload_codec is None:
+            raise ValueError("no payload codec set")
+
+        sig = inspect.signature(invocation_handler)
+
+        params = list(sig.parameters.values())
+        if len(params) > 1:
+            raise ValueError("invocation handler must accept 0 or 1 argument")
+
+        if len(params) == 1:
+            # get parameter's type hint
+            param_type = params[0].annotation
+            if param_type is inspect._empty:
+                raise TypeError("invocation handler parameter must have a type annotation")
+        else:
+            param_type = None
+
+        def _invocation_handler(invocation: types.Invocation):
+            request_obj = None
+            if param_type is not None:
+                if len(invocation.args) != 1:
+                    raise ValueError("only one argument expected in invocation")
+
+                request_obj = self._payload_codec.decode(invocation.args[0], param_type)
+
+            result = invocation_handler(request_obj) if param_type is not None else invocation_handler()
+
+            # no return type in invocation handler
+            if sig.return_annotation is inspect._empty or result is None:
+                return None
+
+            encoded = self._payload_codec.encode(result)
+
+            return types.Result(args=[encoded])
+
+        return self.register(procedure, _invocation_handler)
+
     def call(
         self,
         procedure: str,