Implement equality check for Request class and optimize request removal in Sender class

This commit is contained in:
Jahongir Qurbonov 2025-05-31 13:30:43 +05:00
parent 964ebcddac
commit 59daea32c7
No known key found for this signature in database
GPG Key ID: 256976CED13D5F2D

View File

@ -157,6 +157,9 @@ class Request(Generic[Return]):
state: RequestState state: RequestState
result: Future[Return] result: Future[Return]
def __eq__(self, value: object) -> bool:
return self is value
@dataclass @dataclass
class Sender: class Sender:
@ -422,18 +425,17 @@ class Sender:
req.result.set_exception(result) req.result.set_exception(result)
def _pop_request(self, msg_id: MsgId) -> Optional[Request[object]]: def _pop_request(self, msg_id: MsgId) -> Optional[Request[object]]:
for i, req in enumerate(self._requests): for req in self._requests:
if isinstance(req.state, Serialized) and req.state.msg_id == msg_id: if isinstance(req.state, Serialized) and req.state.msg_id == msg_id:
raise RuntimeError("got response for unsent request") raise RuntimeError("got response for unsent request")
elif isinstance(req.state, Sent) and req.state.msg_id == msg_id: elif isinstance(req.state, Sent) and req.state.msg_id == msg_id:
del self._requests[i] self._requests.remove(req)
return req return req
return None return None
def _drain_requests(self, msg_id: MsgId) -> Iterator[Request[object]]: def _drain_requests(self, msg_id: MsgId) -> Iterator[Request[object]]:
for i in reversed(range(len(self._requests))): for req in self._requests:
req = self._requests[i]
if isinstance(req.state, Serialized) and ( if isinstance(req.state, Serialized) and (
req.state.msg_id == msg_id or req.state.container_msg_id == msg_id req.state.msg_id == msg_id or req.state.container_msg_id == msg_id
): ):
@ -441,7 +443,8 @@ class Sender:
elif isinstance(req.state, Sent) and ( elif isinstance(req.state, Sent) and (
req.state.msg_id == msg_id or req.state.container_msg_id == msg_id req.state.msg_id == msg_id or req.state.container_msg_id == msg_id
): ):
yield self._requests.pop(i) self._requests.remove(req)
yield req
@property @property
def auth_key(self) -> Optional[bytes]: def auth_key(self) -> Optional[bytes]: