From 3041650b4b5b66d41a14a6381d47d8c0828d21a4 Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Sun, 16 Jan 2022 13:17:05 +0800
Subject: [PATCH] :sparkles: add advanced message slice support

---
 nonebot/adapters/_message.py | 48 +++++++++++++++++++++++++++++++++++-
 1 file changed, 47 insertions(+), 1 deletion(-)

diff --git a/nonebot/adapters/_message.py b/nonebot/adapters/_message.py
index 52db50ee..e6bcc665 100644
--- a/nonebot/adapters/_message.py
+++ b/nonebot/adapters/_message.py
@@ -5,12 +5,14 @@ from typing import (
     Any,
     Dict,
     List,
+    Tuple,
     Type,
     Union,
     Generic,
     Mapping,
     TypeVar,
     Iterable,
+    overload,
 )
 
 from ._template import MessageTemplate
@@ -87,7 +89,7 @@ class MessageSegment(Mapping, abc.ABC, Generic[TM]):
         raise NotImplementedError
 
 
-class Message(List[TMS], abc.ABC):
+class Message(List[TMS], Generic[TMS], abc.ABC):
     """消息数组"""
 
     def __init__(
@@ -180,6 +182,50 @@ class Message(List[TMS], abc.ABC):
             self.extend(self._construct(other))
         return self
 
+    @overload
+    def __getitem__(self: TM, __args: str) -> TM:
+        ...
+
+    @overload
+    def __getitem__(self, __args: Tuple[str, int]) -> TMS:
+        ...
+
+    @overload
+    def __getitem__(self: TM, __args: Tuple[str, slice]) -> TM:
+        ...
+
+    @overload
+    def __getitem__(self, __args: int) -> TMS:
+        ...
+
+    @overload
+    def __getitem__(self: TM, __args: slice) -> TM:
+        ...
+
+    def __getitem__(
+        self: TM,
+        __args: Union[
+            str,
+            Tuple[str, int],
+            Tuple[str, slice],
+            int,
+            slice,
+        ],
+    ) -> Union[TMS, TM]:
+        arg1, arg2 = __args if isinstance(__args, tuple) else (__args, None)
+        if isinstance(arg1, int) and arg2 is None:
+            return super().__getitem__(arg1)
+        elif isinstance(arg1, slice) and arg2 is None:
+            return self.__class__(super().__getitem__(arg1))
+        elif isinstance(arg1, str) and arg2 is None:
+            return self.__class__(seg for seg in self if seg.type == arg1)
+        elif isinstance(arg1, str) and isinstance(arg2, int):
+            return [seg for seg in self if seg.type == arg1][arg2]
+        elif isinstance(arg1, str) and isinstance(arg2, slice):
+            return self.__class__([seg for seg in self if seg.type == arg1][arg2])
+        else:
+            raise ValueError("Invalid arguments to __getitem__")
+
     def append(self: TM, obj: Union[str, TMS]) -> TM:
         """
         添加一个消息段到消息数组末尾