Mirrored_Repos/python-kasa
1
0
Fork 0
mirror of https://github.com/python-kasa/python-kasa.git synced 2025-04-26 16:46:23 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

handle cancellation

Browse Source
This commit is contained in:
J. Nick Koston 2024-01-12 15:59:47 -10:00
parent bd285a9fb0
commit e0c614cf62
No known key found for this signature in database
1 changed files with 18 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
kasa/protocol.py
View File

@ -100,14 +100,20 @@ class TPLinkSmartHomeProtocol:
async def close(self) -> None: async def close(self) -> None:
"""Close the connection.""" """Close the connection."""
writer = self.writer writer = self.writer
self.close_without_wait()
if writer:
with contextlib.suppress(Exception):
await writer.wait_closed()
def close_without_wait(self) -> None:
"""Close the connection without waiting for the connection to close."""
writer = self.writer
self.reader = self.writer = None self.reader = self.writer = None
debug_log = _LOGGER.isEnabledFor(logging.DEBUG) debug_log = _LOGGER.isEnabledFor(logging.DEBUG)
if writer: if writer:
if debug_log: if debug_log:
_LOGGER.debug("%s: closing connection", self.host) _LOGGER.debug("%s: closing connection", self.host)
writer.close() writer.close()
with contextlib.suppress(Exception):
await writer.wait_closed()
elif debug_log: elif debug_log:
_LOGGER.debug("%s: connection already closed", self.host) _LOGGER.debug("%s: connection already closed", self.host)
@ -151,6 +157,14 @@ class TPLinkSmartHomeProtocol:
f" {self.host}:{self.port}: {ex}" f" {self.host}:{self.port}: {ex}"
) from ex ) from ex
continue continue
except BaseException as ex:
# Likely something cancelled the task so we need to close the connection
self.close_without_wait()
_LOGGER.debug(
"BaseException during connect, closing connection: %s",
self.host,
ex,
)
try: try:
assert self.reader is not None # noqa: S101 assert self.reader is not None # noqa: S101
@ -169,8 +183,8 @@ class TPLinkSmartHomeProtocol:
"Unable to query the device %s, retrying: %s", self.host, ex "Unable to query the device %s, retrying: %s", self.host, ex
) )
except BaseException as ex: except BaseException as ex:
# Likely something cancelled the task, so we need to close the connection # Likely something cancelled the task so we need to close the connection
await self.close() self.close_without_wait()
_LOGGER.debug( _LOGGER.debug(
"BaseException during query, closing connection: %s", self.host, ex "BaseException during query, closing connection: %s", self.host, ex
) )