(ws) Add network module

Author: protocolstardust

docs/docs/content/documentation/websocket.md

@@ -0,0 +1,86 @@
+# :material-web: WebSocket Server
+
+Rayforce-Py provides WebSocket server functionality that allows you to expose Rayforce queries over WebSocket connections. This enables WebSocket-compatible clients to interact with Rayforce using the standard WebSocket protocol.
+
+## :material-information: Overview
+
+The WebSocket implementation in Rayforce-Py consists of:
+
+- **`WebSocketServer`** - An async WebSocket server that accepts connections and executes Rayforce queries
+- **`WebSocketConnection`** - Represents an individual WebSocket connection (used internally)
+
+The WebSocket server uses Rayforce's native IPC protocol for message serialization, ensuring compatibility with Rayforce's binary format while providing a standard WebSocket interface.
+
+!!! note "Optional Dependency"
+    WebSocket functionality requires the `websockets` library. Install it with:
+    ```bash
+    pip install websockets
+    ```
+
+---
+
+## :material-server: WebSocket Server
+
+The `WebSocketServer` class creates an async WebSocket server that listens for incoming connections and processes Rayforce queries.
+
+### Installation
+
+First, ensure you have the `websockets` library installed:
+
+```bash
+pip install websockets
+```
+
+### Creating a Server
+
+Import `WebSocketServer` from the network module:
+
+```python
+>>> from rayforce.network.websocket import WebSocketServer
+
+>>> server = WebSocketServer(port=8765)
+>>> server
+WebSocketServer(port=8765)
+```
+
+### Starting the Server
+
+The `run()` method starts the server and runs until stopped (e.g., via Ctrl+C):
+
+```python
+>>> import asyncio
+>>> from rayforce.network.websocket import WebSocketServer
+
+>>> async def main():
+...     server = WebSocketServer(port=8765)
+...     await server.run()
+
+>>> asyncio.run(main())
+Rayforce WebSocket Server listening on 0.0.0.0:8765
+^C
+Rayforce WebSocket Server stopped.
+```
+
+!!! note "Async Operation"
+    The `run()` method is **async** and must be called with `await` or within an async context.
+
+---
+
+## :material-code-tags: Message Format
+
+The WebSocket server accepts two types of messages:
+
+### Text Messages
+
+Text messages are treated as Rayforce string queries and are evaluated directly
+
+### Binary Messages
+
+Binary messages must be in Rayforce's native IPC format:
+- **Header**: 16 bytes (prefix, version, flags, endian, msgtype, size)
+- **Payload**: Serialized RayObject using `ser_obj`
+
+The server deserializes the binary data, evaluates it, and returns the result in the same binary format.
+
+
+#### Next: [IPC Documentation](./IPC.md)

docs/mkdocs.yml

@@ -134,6 +134,7 @@ nav:
       - Order By: content/documentation/query-guide/order-by.md
       - Group By: content/documentation/query-guide/group-by.md
     - IPC: content/documentation/IPC.md
+    - Websocket: content/documentation/websocket.md
     - Plugins:
       - Polars: content/documentation/plugins/polars.md
       - Pandas: content/documentation/plugins/pandas.md

rayforce/__init__.py

@@ -64,6 +64,11 @@
     RayforceUserError,
     RayforceValueError,
 )
+from .network import (  # noqa: E402
+    IPCClient,
+    IPCConnection,
+    IPCServer,
+)
 from .types import (  # noqa: E402
     B8,
     C8,
@@ -90,9 +95,6 @@
     Vector,
 )
 from .utils import (  # noqa: E402
-    IPCClient,
-    IPCConnection,
-    IPCServer,
     eval_obj,
     eval_str,
     python_to_ray,

rayforce/network/__init__.py

@@ -0,0 +1,9 @@
+from .ipc import IPCClient, IPCConnection, IPCServer
+
+__all__ = [
+    "IPCClient",
+    "IPCConnection",
+    "IPCServer",
+    # Do not import websockets here to avoid import errors.
+    # Websockets aren't core, but rather an optional plugin.
+]

rayforce/utils/ipc.py rayforce/network/ipc.pyrayforce/utils/ipc.py renamed to rayforce/network/ipc.py

@@ -14,13 +14,10 @@
 from rayforce.types.scalars.numeric.unsigned import U8
 
 try:
-    import asyncio
-
     from websockets import Server, serve  # type: ignore[import-not-found]
 except ImportError as e:
-    module_name = str(e).split("'")[1]
     raise ImportError(
-        f"{module_name} library is required. Install it with: pip install {module_name}"
+        "websockets library is required. Install it with: pip install websockets"
     ) from e
 
 

rayforce/utils/websocket.py rayforce/network/websocket.pyrayforce/utils/websocket.py renamed to rayforce/network/websocket.py

@@ -1,11 +1,7 @@
 from .conversion import python_to_ray, ray_to_python
 from .evaluation import eval_obj, eval_str
-from .ipc import IPCClient, IPCConnection, IPCServer
 
 __all__ = [
-    "IPCClient",
-    "IPCConnection",
-    "IPCServer",
     "eval_obj",
     "eval_str",
     "python_to_ray",

rayforce/utils/__init__.py

@@ -6,7 +6,7 @@
 from rayforce.types import Table
 from rayforce import errors
 from rayforce.ffi import FFI
-from rayforce.utils.ipc import (
+from rayforce.network.ipc import (
     IPCConnection,
     IPCClient,
     IPCServer,
@@ -44,11 +44,11 @@ def mock_handle(self):
 
     @pytest.fixture
     def connection(self, mock_engine, mock_handle):
-        with patch("rayforce.utils.ipc.FFI.get_obj_type", return_value=r.TYPE_I64):
+        with patch("rayforce.network.ipc.FFI.get_obj_type", return_value=r.TYPE_I64):
             return IPCConnection(engine=mock_engine, handle=mock_handle)
 
-    @patch("rayforce.utils.ipc.FFI.write")
-    @patch("rayforce.utils.ipc.ray_to_python")
+    @patch("rayforce.network.ipc.FFI.write")
+    @patch("rayforce.network.ipc.ray_to_python")
     def test_execute(self, mock_ray_to_python, mock_write, connection):
         mock_write.return_value = MagicMock()
         mock_ray_to_python.return_value = "result"
@@ -63,7 +63,7 @@ def test_execute_closed(self, connection):
         with pytest.raises(errors.RayforceIPCError, match="Cannot write to closed connection"):
             connection.execute("test_query")
 
-    @patch("rayforce.utils.ipc.FFI.hclose")
+    @patch("rayforce.network.ipc.FFI.hclose")
     def test_close_removes_from_pool(self, mock_hclose, mock_engine, mock_handle):
         conn = IPCConnection(engine=mock_engine, handle=mock_handle)
         conn_id = id(conn)
@@ -74,7 +74,7 @@ def test_close_removes_from_pool(self, mock_hclose, mock_engine, mock_handle):
         assert conn._closed is True
         mock_hclose.assert_called_once_with(conn.handle)
 
-    @patch("rayforce.utils.ipc.FFI.hclose")
+    @patch("rayforce.network.ipc.FFI.hclose")
     def test_close_idempotent(self, mock_hclose, connection):
         connection.close()
         connection.close()
@@ -92,8 +92,8 @@ class TestIPCEngine:
     def engine(self):
         return IPCClient(host="localhost", port=5000)
 
-    @patch("rayforce.utils.ipc.FFI.get_obj_type")
-    @patch("rayforce.utils.ipc.FFI.hopen")
+    @patch("rayforce.network.ipc.FFI.get_obj_type")
+    @patch("rayforce.network.ipc.FFI.hopen")
     def test_acquire_success(self, mock_hopen, mock_get_obj_type, engine):
         mock_handle = MagicMock(spec=r.RayObject)
 
@@ -111,9 +111,9 @@ def get_obj_type_side_effect(obj):
         assert conn.handle == mock_handle
         assert id(conn) in engine.pool
 
-    @patch("rayforce.utils.ipc.FFI.get_obj_type")
-    @patch("rayforce.utils.ipc.FFI.hopen")
-    @patch("rayforce.utils.ipc.FFI.get_error_obj")
+    @patch("rayforce.network.ipc.FFI.get_obj_type")
+    @patch("rayforce.network.ipc.FFI.hopen")
+    @patch("rayforce.network.ipc.FFI.get_error_obj")
     def test_acquire_failure(self, mock_get_error, mock_hopen, mock_get_obj_type, engine):
         mock_error = MagicMock(spec=r.RayObject)
 
@@ -129,8 +129,8 @@ def get_obj_type_side_effect(obj):
         with pytest.raises(errors.RayforceIPCError, match="Error when establishing connection"):
             engine.acquire()
 
-    @patch("rayforce.utils.ipc.FFI.get_obj_type")
-    @patch("rayforce.utils.ipc.FFI.hopen")
+    @patch("rayforce.network.ipc.FFI.get_obj_type")
+    @patch("rayforce.network.ipc.FFI.hopen")
     def test_dispose_connections(self, mock_hopen, mock_get_obj_type, engine):
         mock_handle1 = MagicMock(spec=r.RayObject)
         mock_handle2 = MagicMock(spec=r.RayObject)
@@ -174,8 +174,8 @@ def test_init_invalid_port_too_low(self, port):
         with pytest.raises(errors.RayforceIPCError, match="Invalid port number"):
             IPCServer(port=port)
 
-    @patch("rayforce.utils.ipc.FFI.ipc_listen")
-    @patch("rayforce.utils.ipc.FFI.runtime_run")
+    @patch("rayforce.network.ipc.FFI.ipc_listen")
+    @patch("rayforce.network.ipc.FFI.runtime_run")
     def test_listen_success(self, mock_runtime_run, mock_ipc_listen, server):
         mock_ipc_listen.return_value = 123
         mock_runtime_run.return_value = 0
@@ -186,9 +186,9 @@ def test_listen_success(self, mock_runtime_run, mock_ipc_listen, server):
         assert server._listener_id == 123
         mock_runtime_run.assert_called_once()
 
-    @patch("rayforce.utils.ipc.FFI.ipc_listen")
-    @patch("rayforce.utils.ipc.FFI.runtime_run")
-    @patch("rayforce.utils.ipc.FFI.ipc_close_listener")
+    @patch("rayforce.network.ipc.FFI.ipc_listen")
+    @patch("rayforce.network.ipc.FFI.runtime_run")
+    @patch("rayforce.network.ipc.FFI.ipc_close_listener")
     def test_listen_closes_on_exception(
         self, mock_close, mock_runtime_run, mock_ipc_listen, server
     ):
@@ -202,9 +202,9 @@ def test_listen_closes_on_exception(
         mock_close.assert_called_once_with(123)
         assert server._listener_id is None
 
-    @patch("rayforce.utils.ipc.FFI.ipc_listen")
-    @patch("rayforce.utils.ipc.FFI.runtime_run")
-    @patch("rayforce.utils.ipc.FFI.ipc_close_listener")
+    @patch("rayforce.network.ipc.FFI.ipc_listen")
+    @patch("rayforce.network.ipc.FFI.runtime_run")
+    @patch("rayforce.network.ipc.FFI.ipc_close_listener")
     def test_listen_closes_on_keyboard_interrupt(
         self, mock_close, mock_runtime_run, mock_ipc_listen, server
     ):
@@ -238,7 +238,7 @@ def mock_handle(self):
 
     @pytest.fixture
     def connection(self, mock_engine, mock_handle):
-        with patch("rayforce.utils.ipc.FFI.get_obj_type", return_value=r.TYPE_I64):
+        with patch("rayforce.network.ipc.FFI.get_obj_type", return_value=r.TYPE_I64):
             return IPCConnection(engine=mock_engine, handle=mock_handle)
 
     def _capture_and_eval(self, connection, query_obj):
@@ -250,8 +250,8 @@ def capture_write(handle, data):
             mock_result = MagicMock(spec=r.RayObject)
             return mock_result
 
-        with patch("rayforce.utils.ipc.FFI.write", side_effect=capture_write):
-            with patch("rayforce.utils.ipc.ray_to_python", return_value="mocked_result"):
+        with patch("rayforce.network.ipc.FFI.write", side_effect=capture_write):
+            with patch("rayforce.network.ipc.ray_to_python", return_value="mocked_result"):
                 connection.execute(query_obj)
 
         assert captured_obj is not None