Coverage for nexios\websockets\channels.py: 31%

1 import sys

2 import os

3 import time

4 import uuid

5 from nexios import logging as nexios_logger

6 from .utils import (

7 ChannelAddStatusEnum,

8 ChannelRemoveStatusEnum,

9 GroupSendStatusEnum,

10 PayloadTypeEnum,

11 ChannelMessageDC,

12)

13 import typing

14 from nexios.websockets import WebSocket

15

16 logging = nexios_logger.getLogger("nexios")

17

18

19 class Channel:

20 def __init__(

21 self,

22 websocket: WebSocket,

23 payload_type: str,

24 expires: typing.Optional[int] = None,

25 ) -> None:

26 """Main websocket channel class.

27

28 Args:

29 websocket (WebSocket): Starlette websocket

30 expires (int): Channel ttl in seconds

31 encoding (str): encoding of payload (str, bytes, json)

32 uuid (str): channel uuid

33 created (tim): channel creation time

34 """

35 assert isinstance(websocket, WebSocket)

36 assert isinstance(expires, int)

37 assert isinstance(payload_type, str) and payload_type in [

38 PayloadTypeEnum.JSON.value,

39 PayloadTypeEnum.TEXT.value,

40 PayloadTypeEnum.BYTES.value,

41 ]

42

43 self.websocket = websocket

44 self.expires = expires

45 self.payload_type = payload_type

46 self.uuid = uuid.uuid4()

47 self.created = time.time()

48

49 async def _send(self, payload: typing.Any) -> None:

50 try:

51 if self.payload_type == "json":

52 await self.websocket.send_json(payload)

53 elif self.payload_type == "text":

54 await self.websocket.send_text(payload)

55 elif self.payload_type == "bytes":

56 await self.websocket.send_bytes(payload)

57 else:

58 await self.websocket.send(payload)

59 except RuntimeError as error:

60 logging.debug(error)

61

62 self.created = time.time()

63

64 async def _is_expired(self) -> bool:

65 if not self.expires:

66 return False

67 return (self.expires + int(self.created)) < time.time()

68

69 def __repr__(self) -> str:

70 return f"{self.__class__.__name__} {self.uuid=} {self.payload_type=} {self.expires=}"

71

72

73 class ChannelBox:

74 CHANNEL_GROUPS: typing.Dict[str, typing.Any] = (

75 {}

76 ) # groups of channels ~ key: group_name, val: dict of channels

77 CHANNEL_GROUPS_HISTORY: typing.Dict[str, typing.Any] = {} # history messages

78 HISTORY_SIZE: int = int(os.getenv("CHANNEL_BOX_HISTORY_SIZE", 1_048_576))

79

80 @classmethod

81 async def add_channel_to_group(

82 cls,

83 channel: Channel,

84 group_name: str = "default",

85 ) -> ChannelAddStatusEnum:

86 """Add channel to group.

87

88 Args:

89 channel (Channel): Instance of Channel class

90 group_name (str): Group name

91

92 """

93 assert group_name, "Group name must to be set."

94

95 if group_name not in cls.CHANNEL_GROUPS:

96 cls.CHANNEL_GROUPS[group_name] = {}

97 channel_add_status = ChannelAddStatusEnum.CHANNEL_ADDED

98 else:

99 channel_add_status = ChannelAddStatusEnum.CHANNEL_EXIST

100

101 cls.CHANNEL_GROUPS[group_name][channel] = ...

102 return channel_add_status

103

104 @classmethod

105 async def remove_channel_from_group(

106 cls,

107 channel: Channel,

108 group_name: str,

109 ) -> ChannelRemoveStatusEnum:

110 """Remove channel from group.

111

112 Args:

113 channel (Channel): Instance of Channel class

114 group_name (str): Group name

115 """

116 channel_remove_status: typing.Any = None

117 if channel in cls.CHANNEL_GROUPS.get(group_name, {}):

118 try:

119 del cls.CHANNEL_GROUPS[group_name][channel]

120 channel_remove_status = ChannelRemoveStatusEnum.CHANNEL_REMOVED

121 except KeyError:

122 channel_remove_status = ChannelRemoveStatusEnum.CHANNEL_DOES_NOT_EXIST

123

124 if not any(cls.CHANNEL_GROUPS.get(group_name, {})):

125 try:

126 del cls.CHANNEL_GROUPS[group_name]

127 channel_remove_status = ChannelRemoveStatusEnum.GROUP_REMOVED

128 except KeyError:

129 channel_remove_status = ChannelRemoveStatusEnum.GROUP_DOES_NOT_EXIST

130

131 await cls._clean_expired()

132 return channel_remove_status

133

134 @classmethod

135 async def group_send(

136 cls,

137 group_name: str = "default",

138 payload: typing.Union[typing.Dict[str, typing.Any], str, bytes] = {},

139 save_history: bool = False,

140 ) -> GroupSendStatusEnum:

141 """Send payload to all channels connected to group.

142

143 Args:

144 group_name (str, optional): Group name

145 payload (dict, optional): Payload to channel

146 save_history (bool, optional): Save message history. Defaults to False.

147

148 """

149 assert group_name, "Group name must to be set."

150

151 if save_history:

152 cls.CHANNEL_GROUPS_HISTORY.setdefault(group_name, [])

153 cls.CHANNEL_GROUPS_HISTORY[group_name].append(

154 ChannelMessageDC(

155 payload=payload, # type:ignore

156 )

157 )

158 if sys.getsizeof(cls.CHANNEL_GROUPS_HISTORY[group_name]) > cls.HISTORY_SIZE:

159 cls.CHANNEL_GROUPS_HISTORY[group_name] = []

160

161 group_send_status = GroupSendStatusEnum.NO_SUCH_GROUP

162 for channel in cls.CHANNEL_GROUPS.get(group_name, {}):

163 await channel._send(payload)

164 group_send_status = GroupSendStatusEnum.GROUP_SEND

165

166 return group_send_status

167

168 @classmethod

169 async def show_groups(cls) -> typing.Dict[str, typing.Any]:

170 return cls.CHANNEL_GROUPS

171

172 @classmethod

173 async def flush_groups(cls) -> None:

174 cls.CHANNEL_GROUPS = {}

175

176 @classmethod

177 async def show_history(

178 cls,

179 group_name: str = "",

180 ) -> typing.Dict[str, typing.Any]:

181 return (

182 cls.CHANNEL_GROUPS_HISTORY.get(group_name, {})

183 if group_name

184 else cls.CHANNEL_GROUPS_HISTORY

185 )

186

187 @classmethod

188 async def flush_history(cls) -> None:

189 cls.CHANNEL_GROUPS_HISTORY = {}

190

191 @classmethod

192 async def _clean_expired(cls) -> None:

193 for group_name in list(cls.CHANNEL_GROUPS):

194 for channel in cls.CHANNEL_GROUPS.get(group_name, {}):

195 _is_expired = await channel._is_expired()

196 if _is_expired:

197 try:

198 del cls.CHANNEL_GROUPS[group_name][channel]

199 except KeyError:

200 logging.debug("No such channel")

201

202 if not any(cls.CHANNEL_GROUPS.get(group_name, {})):

203 try:

204 del cls.CHANNEL_GROUPS[group_name]

205 except KeyError:

206 logging.debug("No such group")

207

208 @classmethod

209 async def close_all_connections(cls) -> None:

210 """

211 Close all WebSocket connections in all groups.

212 """

213 for group_name, channels in cls.CHANNEL_GROUPS.items():

214 for channel in list(

215 channels.keys()

216 ): # Use list() to avoid RuntimeError due to dict size change

217 try:

218 await channel.websocket.close()

219 logging.debug(

220 f"Closed connection for channel {channel.uuid} in group {group_name}"

221 )

222 except Exception as e:

223 logging.error(

224 f"Failed to close connection for channel {channel.uuid}: {e}"

225 )

226

227 cls.CHANNEL_GROUPS = {}

228 logging.debug("All connections closed and groups cleared.")