[OneBot] Fix v11 custom segment init and resolve errors

Meloland · Dec 11, 2024 · f006ee0 · f006ee0
1 parent 4bddb6a
commit f006ee0
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 7 deletions.
diff --git a/src/melobot/protocols/onebot/v11/adapter/segment.py b/src/melobot/protocols/onebot/v11/adapter/segment.py
@@ -308,13 +308,14 @@ def add_type(
 
         seg_cls = type(
             type_classname,
-            (Segment,),
+            (_CustomSegInterface,),
             {
                 "Model": create_model(
                     type_dataname,
                     type=(seg_type_hint, ...),
                     data=(seg_data_hint, ...),
-                )
+                ),
+                "SegTypeVal": type_name,
             },
         )
         setattr(
@@ -342,7 +343,10 @@ def data(self) -> _SegDataT:
     @classmethod
     def resolve(cls, seg_type: Any, seg_data: Any) -> Segment:
         cls_name = f"{seg_type.lower().capitalize()}Segment"
-        cls_map = {subcls.__name__: subcls for subcls in cls.__subclasses__()}
+        cls_map = {
+            subcls.__name__: subcls
+            for subcls in cls.__subclasses__() + _CustomSegInterface.__subclasses__()
+        }
         if cls_name in cls_map:
             return cls_map[cls_name].resolve(seg_type, seg_data)
         return cls(seg_type, **seg_data)
@@ -370,9 +374,13 @@ def to_json(self, force_str: bool = False) -> str:
 
 
 class _CustomSegInterface(Segment[_SegTypeT, _SegDataT]):
-    def __init__(  # pylint: disable=super-init-not-called,unused-argument
-        self, **data: Any
-    ) -> None: ...
+    SegTypeVal: str
+
+    def __init__(self, seg_type: _SegTypeT | None = None, **seg_data: _SegDataT) -> None:
+        if seg_type is None:
+            super().__init__(cast(_SegTypeT, self.__class__.SegTypeVal), **seg_data)
+        else:
+            super().__init__(seg_type, **seg_data)
 
 
 class _TextData(TypedDict):

diff --git a/tests/onebot/v11/test_adapter_segment.py b/tests/onebot/v11/test_adapter_segment.py
@@ -31,7 +31,14 @@ async def test_add_type():
     s = SType(key="123")
     assert s.type == "MyS"
     assert s.data == {"key": "123"}
-    assert SType in seg.Segment.__subclasses__()
+    assert s.raw == {"type": "MyS", "data": {"key": "123"}}
+
+    s2 = seg.Segment.resolve("MyS", {"key": "123"})
+    assert s.type == "MyS"
+    assert s.data == {"key": "123"}
+    assert s.raw == {"type": "MyS", "data": {"key": "123"}}
+    assert isinstance(s2, SType)
+
     with pt.raises(ValidationError):
         SType(key=123)