@@ -72,24 +72,15 @@ async def send(self, socket: str, connection_id: str, data: bytes):
7272 raise exception_from_grpc_error (grpc_err ) from grpc_err
7373
7474
75- class WebsocketWorkerOptions :
76- """Options for websocket workers."""
77-
78- def __init__ (self , socket_name : str , event_type : Literal ["connect" , "disconnect" , "message" ]):
79- """Construct new websocket worker options."""
80- self .socket_name = socket_name
81- self .event_type = WebsocketWorkerOptions ._to_grpc_event_type (event_type )
82-
83- @staticmethod
84- def _to_grpc_event_type (event_type : Literal ["connect" , "disconnect" , "message" ]) -> WebsocketEventType :
85- if event_type == "connect" :
86- return WebsocketEventType .Connect
87- elif event_type == "disconnect" :
88- return WebsocketEventType .Disconnect
89- elif event_type == "message" :
90- return WebsocketEventType .Message
91- else :
92- raise ValueError (f"Event type { event_type } is unsupported" )
75+ def _to_grpc_event_type (event_type : Literal ["connect" , "disconnect" , "message" ]) -> WebsocketEventType :
76+ if event_type == "connect" :
77+ return WebsocketEventType .Connect
78+ elif event_type == "disconnect" :
79+ return WebsocketEventType .Disconnect
80+ elif event_type == "message" :
81+ return WebsocketEventType .Message
82+ else :
83+ raise ValueError (f"Event type { event_type } is unsupported" )
9384
9485
9586def _to_resource (b : Websocket ) -> ResourceIdentifier :
@@ -194,11 +185,15 @@ class WebsocketWorker(FunctionServer):
194185 _registration_request : RegistrationRequest
195186 _responses : AsyncNotifierList [ClientMessage ]
196187
197- def __init__ (self , socket_name : str , event_type : WebsocketEventType , handler : WebsocketHandler ):
188+ def __init__ (
189+ self , socket_name : str , event_type : Literal ["connect" , "disconnect" , "message" ], handler : WebsocketHandler
190+ ):
198191 """Construct a new WebsocketHandler."""
199192 self ._handler = handler
200193 self ._responses = AsyncNotifierList ()
201- self ._registration_request = RegistrationRequest (socket_name = socket_name , event_type = event_type )
194+ self ._registration_request = RegistrationRequest (
195+ socket_name = socket_name , event_type = _to_grpc_event_type (event_type )
196+ )
202197
203198 Nitric ._register_worker (self )
204199
0 commit comments