Add api registry

python/hypha_rpc/__init__.py

@@ -22,18 +22,22 @@ class API(ObjectProxy):
  def __init__(self, *args, **kwargs):
  super().__init__(*args, **kwargs)
  self._registry = {}
+ self._export_handler = self._default_export_handler
 
- def list(self):
- return list(self._registry.keys())
-
- def get(self, name):
- return self._registry.get(name)
-
- def export(self, obj, config=None):
+ def _default_export_handler(self, obj, config=None):
  config = config or {}
  name = config.get("name", shortuuid.uuid())
  self._registry[name] = obj
 
+ def set_export_handler(self, handler):
+ self._export_handler = handler
+
+ def export(self, obj, config=None):
+ return self._export_handler(obj, config=config)
+
+ def get_registry(self):
+ return self._registry
+
 
 # An placeholder object for the API
 api = API()

python/tests/test_websocket_rpc.py

@@ -201,13 +201,25 @@ 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}, {"name": "hello"})
+ assert "hello" in api.get_registry()
+
+ api.get_registry().clear()
+ api.set_export_handler(lambda obj, config=None: None)
  api.export({"hello": lambda x: "hello " + x})
+ assert "hello" not in api.get_registry().keys()
+
+ api.get_registry().clear()
+ api.set_export_handler(api._default_export_handler)
  api.export({"hello": lambda x: "hello " + x}, {"name": "hello2"})
- assert "hello2" in api.list()
+ assert "hello2" in api.get_registry().keys()
+
 
 @pytest.mark.asyncio
 async def test_connect_to_server(websocket_server):