Implement .count and optimize .get performance for message slice

This commit is contained in:
Mix 2022-01-17 00:28:36 +08:00
parent 1221baaa94
commit 3b4c4d3081

View File

@ -12,6 +12,7 @@ from typing import (
Mapping, Mapping,
TypeVar, TypeVar,
Iterable, Iterable,
Optional,
overload, overload,
) )
@ -232,14 +233,20 @@ class Message(List[TMS], abc.ABC):
return super().index(first_segment, *args) # type: ignore return super().index(first_segment, *args) # type: ignore
return super().index(value, *args) return super().index(value, *args)
def get(self: TM, type_: str, count: int) -> TM: def get(self: TM, type_: str, count: Optional[int] = None) -> TM:
iterator = (seg for seg in self if seg.type == type_) if count is None:
return self.__class__( return self[type_]
filter(
lambda seg: seg is not None, iterator, filtered = (seg for seg in self if seg.type == type_), []
(next(iterator) for _ in range(count)), for _ in range(count):
) seg = next(iterator, None)
) if seg is None:
break
filtered.append(seg)
return self.__class__(filtered)
def count(self, value: Union[TMS, str]) -> int:
return len(self[value]) if isinstance(value, str) else super().count(value)
def append(self: TM, obj: Union[str, TMS]) -> TM: def append(self: TM, obj: Union[str, TMS]) -> TM:
""" """