155 lines
4.2 KiB
Python
155 lines
4.2 KiB
Python
from typing import Any, Dict, List, Type, TypeVar, Union
|
|
|
|
import attr
|
|
from pydantic import BaseModel, GetCoreSchemaHandler
|
|
from pydantic_core import CoreSchema, core_schema
|
|
|
|
from ..models.client_metrics import ClientMetrics
|
|
from ..models.modeling_cmd import ModelingCmd
|
|
from ..models.modeling_cmd_id import ModelingCmdId
|
|
from ..models.modeling_cmd_req import ModelingCmdReq
|
|
from ..models.rtc_ice_candidate_init import RtcIceCandidateInit
|
|
from ..models.rtc_session_description import RtcSessionDescription
|
|
|
|
|
|
class trickle_ice(BaseModel):
|
|
"""The trickle ICE candidate request."""
|
|
|
|
candidate: RtcIceCandidateInit
|
|
|
|
type: str = "trickle_ice"
|
|
|
|
|
|
class sdp_offer(BaseModel):
|
|
"""The SDP offer request."""
|
|
|
|
offer: RtcSessionDescription
|
|
|
|
type: str = "sdp_offer"
|
|
|
|
|
|
class modeling_cmd_req(BaseModel):
|
|
"""The modeling command request."""
|
|
|
|
cmd: ModelingCmd
|
|
|
|
cmd_id: ModelingCmdId
|
|
|
|
type: str = "modeling_cmd_req"
|
|
|
|
|
|
class modeling_cmd_batch_req(BaseModel):
|
|
"""A sequence of modeling requests. If any request fails, following requests will not be tried."""
|
|
|
|
requests: List[ModelingCmdReq]
|
|
|
|
type: str = "modeling_cmd_batch_req"
|
|
|
|
|
|
class ping(BaseModel):
|
|
"""The client-to-server Ping to ensure the WebSocket stays alive."""
|
|
|
|
type: str = "ping"
|
|
|
|
|
|
class metrics_response(BaseModel):
|
|
"""The response to a metrics collection request from the server."""
|
|
|
|
metrics: ClientMetrics
|
|
|
|
type: str = "metrics_response"
|
|
|
|
|
|
GY = TypeVar("GY", bound="WebSocketRequest")
|
|
|
|
|
|
@attr.s(auto_attribs=True)
|
|
class WebSocketRequest:
|
|
|
|
"""The websocket messages the server receives."""
|
|
|
|
type: Union[
|
|
trickle_ice,
|
|
sdp_offer,
|
|
modeling_cmd_req,
|
|
modeling_cmd_batch_req,
|
|
ping,
|
|
metrics_response,
|
|
]
|
|
|
|
def __init__(
|
|
self,
|
|
type: Union[
|
|
trickle_ice,
|
|
sdp_offer,
|
|
modeling_cmd_req,
|
|
modeling_cmd_batch_req,
|
|
ping,
|
|
metrics_response,
|
|
],
|
|
):
|
|
self.type = type
|
|
|
|
def model_dump(self) -> Dict[str, Any]:
|
|
if isinstance(self.type, trickle_ice):
|
|
WI: trickle_ice = self.type
|
|
return WI.model_dump()
|
|
elif isinstance(self.type, sdp_offer):
|
|
YR: sdp_offer = self.type
|
|
return YR.model_dump()
|
|
elif isinstance(self.type, modeling_cmd_req):
|
|
XK: modeling_cmd_req = self.type
|
|
return XK.model_dump()
|
|
elif isinstance(self.type, modeling_cmd_batch_req):
|
|
OB: modeling_cmd_batch_req = self.type
|
|
return OB.model_dump()
|
|
elif isinstance(self.type, ping):
|
|
QQ: ping = self.type
|
|
return QQ.model_dump()
|
|
elif isinstance(self.type, metrics_response):
|
|
WX: metrics_response = self.type
|
|
return WX.model_dump()
|
|
|
|
raise Exception("Unknown type")
|
|
|
|
@classmethod
|
|
def from_dict(cls: Type[GY], d: Dict[str, Any]) -> GY:
|
|
if d.get("type") == "trickle_ice":
|
|
QL: trickle_ice = trickle_ice(**d)
|
|
return cls(type=QL)
|
|
elif d.get("type") == "sdp_offer":
|
|
ME: sdp_offer = sdp_offer(**d)
|
|
return cls(type=ME)
|
|
elif d.get("type") == "modeling_cmd_req":
|
|
EB: modeling_cmd_req = modeling_cmd_req(**d)
|
|
return cls(type=EB)
|
|
elif d.get("type") == "modeling_cmd_batch_req":
|
|
VK: modeling_cmd_batch_req = modeling_cmd_batch_req(**d)
|
|
return cls(type=VK)
|
|
elif d.get("type") == "ping":
|
|
ZC: ping = ping(**d)
|
|
return cls(type=ZC)
|
|
elif d.get("type") == "metrics_response":
|
|
BE: metrics_response = metrics_response(**d)
|
|
return cls(type=BE)
|
|
|
|
raise Exception("Unknown type")
|
|
|
|
@classmethod
|
|
def __get_pydantic_core_schema__(
|
|
cls, source_type: Any, handler: GetCoreSchemaHandler
|
|
) -> CoreSchema:
|
|
return core_schema.no_info_after_validator_function(
|
|
cls,
|
|
handler(
|
|
Union[
|
|
trickle_ice,
|
|
sdp_offer,
|
|
modeling_cmd_req,
|
|
modeling_cmd_batch_req,
|
|
ping,
|
|
metrics_response,
|
|
]
|
|
),
|
|
)
|