add websocket class and coolq message segment

nonebot/adapters/__init__.py

@@ -21,22 +21,46 @@ class BaseBot(object):
 class BaseMessageSegment(dict):
 
     def __init__(self,
-                 d: Optional[Dict[str, Any]] = None,
-                 *,
                  type_: Optional[str] = None,
                  data: Optional[Dict[str, str]] = None):
         super().__init__()
-        if isinstance(d, dict) and d.get('type'):
-            self.update(d)
-        elif type_:
+        if type_:
             self.type = type_
             self.data = data
         else:
-            raise ValueError('the "type" field cannot be None or empty')
+            raise ValueError('The "type" field cannot be empty')
 
     def __str__(self):
         raise NotImplementedError
 
+    def __getitem__(self, item):
+        if item not in ("type", "data"):
+            raise KeyError(f'Key "{item}" is not allowed')
+        return super().__getitem__(item)
+
+    def __setitem__(self, key, value):
+        if key not in ("type", "data"):
+            raise KeyError(f'Key "{key}" is not allowed')
+        return super().__setitem__(key, value)
+
+    # TODO: __eq__ __add__
+
+    @property
+    def type(self) -> str:
+        return self["type"]
+
+    @type.setter
+    def type(self, value: str):
+        self["type"] = value
+
+    @property
+    def data(self) -> Dict[str, str]:
+        return self["data"]
+
+    @data.setter
+    def data(self, data: Optional[Dict[str, str]]):
+        self["data"] = data or {}
+
 
 class BaseMessage(list):
 

nonebot/adapters/coolq.py

@@ -5,13 +5,40 @@ import httpx
 
 from nonebot.event import Event
 from nonebot.config import Config
-from nonebot.adapters import BaseBot, BaseMessage, BaseMessageSegment
 from nonebot.message import handle_event
+from nonebot.drivers import BaseWebSocket
+from nonebot.adapters import BaseBot, BaseMessage, BaseMessageSegment
+
+
+def escape(s: str, *, escape_comma: bool = True) -> str:
+    """
+    对字符串进行 CQ 码转义。
+
+    ``escape_comma`` 参数控制是否转义逗号（``,``）。
+    """
+    s = s.replace("&", "&") \
+        .replace("[", "[") \
+        .replace("]", "]")
+    if escape_comma:
+        s = s.replace(",", ",")
+    return s
+
+
+def unescape(s: str) -> str:
+    """对字符串进行 CQ 码去转义。"""
+    return s.replace(",", ",") \
+        .replace("[", "[") \
+        .replace("]", "]") \
+        .replace("&", "&")
 
 
 class Bot(BaseBot):
 
-    def __init__(self, type_: str, config: Config, *, websocket=None):
+    def __init__(self,
+                 type_: str,
+                 config: Config,
+                 *,
+                 websocket: BaseWebSocket = None):
         if type_ not in ["http", "websocket"]:
             raise ValueError("Unsupported connection type")
         self.type = type_
@@ -33,7 +60,32 @@ class Bot(BaseBot):
 
 
 class MessageSegment(BaseMessageSegment):
-    pass
+
+    def __str__(self):
+        type_ = self.type
+        data = self.data.copy()
+
+        # process special types
+        if type_ == "text":
+            return escape(data.get("text", ""), escape_comma=False)
+        elif type_ == "at_all":
+            type_ = "at"
+            data = {"qq": "all"}
+
+        params = ",".join([f"{k}={escape(str(v))}" for k, v in data.items()])
+        return f"[CQ:{type_}{',' if params else ''}{params}]"
+
+    @staticmethod
+    def at(user_id: int) -> "MessageSegment":
+        return MessageSegment("at", {"qq": str(user_id)})
+
+    @staticmethod
+    def at_all() -> "MessageSegment":
+        return MessageSegment("at_all")
+
+    @staticmethod
+    def dice() -> "MessageSegment":
+        return MessageSegment(type_="dice")
 
 
 class Message(BaseMessage):