FIXME: Make TProtocol Optional[TProtocol] to keep types consistent (#770)

* FIXME: Make TProtocol Optional[TProtocol] to keep types consistent

* correct test case of test_protocol_muxer

* add newsfragment

* unit test added

---------

Co-authored-by: Manu Sheel Gupta <manusheel.edu@gmail.com>
This commit is contained in:
Jinesh Jain
2025-08-20 06:50:37 +05:30
committed by GitHub
parent e20a9a3814
commit dabb3a0962
9 changed files with 167 additions and 17 deletions

View File

@ -48,12 +48,11 @@ class Multiselect(IMultiselectMuxer):
"""
self.handlers[protocol] = handler
# FIXME: Make TProtocol Optional[TProtocol] to keep types consistent
async def negotiate(
self,
communicator: IMultiselectCommunicator,
negotiate_timeout: int = DEFAULT_NEGOTIATE_TIMEOUT,
) -> tuple[TProtocol, StreamHandlerFn | None]:
) -> tuple[TProtocol | None, StreamHandlerFn | None]:
"""
Negotiate performs protocol selection.
@ -84,14 +83,14 @@ class Multiselect(IMultiselectMuxer):
raise MultiselectError() from error
else:
protocol = TProtocol(command)
if protocol in self.handlers:
protocol_to_check = None if not command else TProtocol(command)
if protocol_to_check in self.handlers:
try:
await communicator.write(protocol)
await communicator.write(command)
except MultiselectCommunicatorError as error:
raise MultiselectError() from error
return protocol, self.handlers[protocol]
return protocol_to_check, self.handlers[protocol_to_check]
try:
await communicator.write(PROTOCOL_NOT_FOUND_MSG)
except MultiselectCommunicatorError as error:

View File

@ -134,8 +134,10 @@ class MultiselectClient(IMultiselectClient):
:raise MultiselectClientError: raised when protocol negotiation failed
:return: selected protocol
"""
# Represent `None` protocol as an empty string.
protocol_str = protocol if protocol is not None else ""
try:
await communicator.write(protocol)
await communicator.write(protocol_str)
except MultiselectCommunicatorError as error:
raise MultiselectClientError() from error
@ -145,7 +147,7 @@ class MultiselectClient(IMultiselectClient):
except MultiselectCommunicatorError as error:
raise MultiselectClientError() from error
if response == protocol:
if response == protocol_str:
return protocol
if response == PROTOCOL_NOT_FOUND_MSG:
raise MultiselectClientError("protocol not supported")

View File

@ -30,7 +30,10 @@ class MultiselectCommunicator(IMultiselectCommunicator):
"""
:raise MultiselectCommunicatorError: raised when failed to write to underlying reader
""" # noqa: E501
msg_bytes = encode_delim(msg_str.encode())
if msg_str is None:
msg_bytes = encode_delim(b"")
else:
msg_bytes = encode_delim(msg_str.encode())
try:
await self.read_writer.write(msg_bytes)
except IOException as error: