From 252e6cc7a8fb353a77e861ce69a9826e9e82e168 Mon Sep 17 00:00:00 2001
From: Wei Ouyang <oeway007@gmail.com>
Date: Tue, 20 Aug 2024 23:45:45 -0700
Subject: [PATCH] Add api.export

---
 python/hypha_rpc/__init__.py       | 25 +++++++++++++++++++++++++
 python/tests/test_websocket_rpc.py |  7 +++++++
 2 files changed, 32 insertions(+)

diff --git a/python/hypha_rpc/__init__.py b/python/hypha_rpc/__init__.py
index 5e58589..b2bb63b 100644
--- a/python/hypha_rpc/__init__.py
+++ b/python/hypha_rpc/__init__.py
@@ -1,6 +1,9 @@
 """Provide hypha-rpc to connecting to Hypha server."""
 
+import shortuuid
+
 from .rpc import RPC
+from .utils import ObjectProxy
 from .sync import connect_to_server as connect_to_server_sync
 from .sync import get_remote_service as get_remote_service_sync
 from .sync import get_rtc_service as get_rtc_service_sync
@@ -14,7 +17,29 @@
     setup_local_client,
 )
 
+
+class API(ObjectProxy):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._registry = {}
+
+    def list(self):
+        return list(self._registry.keys())
+
+    def get(self, name):
+        return self._registry.get(name)
+
+    def export(self, obj, config=None):
+        config = config or {}
+        name = config.get("name", shortuuid.uuid())
+        self._registry[name] = obj
+
+
+# An placeholder object for the API
+api = API()
+
 __all__ = [
+    "api",
     "RPC",
     "login",
     "connect_to_server",
diff --git a/python/tests/test_websocket_rpc.py b/python/tests/test_websocket_rpc.py
index b544e72..dbb5a7f 100644
--- a/python/tests/test_websocket_rpc.py
+++ b/python/tests/test_websocket_rpc.py
@@ -201,6 +201,13 @@ def test_connect_to_server_sync(websocket_server):
         info = login.start()
         assert "key" in info
 
+@pytest.mark.asyncio
+async def test_export_api(websocket_server):
+    """Test exporting API."""
+    from hypha_rpc import api
+    api.export({"hello": lambda x: "hello " + x})
+    api.export({"hello": lambda x: "hello " + x}, {"name": "hello2"})
+    assert "hello2" in api.list()
 
 @pytest.mark.asyncio
 async def test_connect_to_server(websocket_server):