Skip to content

Commit ce02ea2

Browse files
skbehskbeh
committed
refactor: use 'with' statement for socket handling
Use a 'with' statement to handle socket creation and cleanup. This ensures that sockets are always closed properly, even when exceptions occur, preventing potential resource leaks.
1 parent c5cc0f2 commit ce02ea2

File tree

2 files changed

+98
-97
lines changed

2 files changed

+98
-97
lines changed

‎miio/miioprotocol.py‎

Lines changed: 91 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -116,34 +116,34 @@ def discover(addr: Optional[str] = None, timeout: int = 5) -> Any:
116116
"21310020ffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
117117
)
118118

119-
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
120-
s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
121-
s.settimeout(timeout)
122-
for _ in range(3):
123-
s.sendto(helobytes, (addr, 54321))
124-
while True:
125-
try:
126-
data, recv_addr = s.recvfrom(1024)
127-
m: Message = Message.parse(data)
128-
_LOGGER.debug("Got a response: %s", m)
129-
if not is_broadcast:
130-
return m
131-
132-
if recv_addr[0] not in seen_addrs:
133-
_LOGGER.info(
134-
" IP %s (ID: %s) - token: %s",
135-
recv_addr[0],
136-
binascii.hexlify(m.header.value.device_id).decode(),
137-
codecs.encode(m.checksum, "hex"),
138-
)
139-
seen_addrs.append(recv_addr[0])
140-
except socket.timeout:
141-
if is_broadcast:
142-
_LOGGER.info("Discovery done")
143-
return # ignore timeouts on discover
144-
except Exception as ex:
145-
_LOGGER.warning("error while reading discover results: %s", ex)
146-
break
119+
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
120+
s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
121+
s.settimeout(timeout)
122+
for _ in range(3):
123+
s.sendto(helobytes, (addr, 54321))
124+
while True:
125+
try:
126+
data, recv_addr = s.recvfrom(1024)
127+
m: Message = Message.parse(data)
128+
_LOGGER.debug("Got a response: %s", m)
129+
if not is_broadcast:
130+
return m
131+
132+
if recv_addr[0] not in seen_addrs:
133+
_LOGGER.info(
134+
" IP %s (ID: %s) - token: %s",
135+
recv_addr[0],
136+
binascii.hexlify(m.header.value.device_id).decode(),
137+
codecs.encode(m.checksum, "hex"),
138+
)
139+
seen_addrs.append(recv_addr[0])
140+
except socket.timeout:
141+
if is_broadcast:
142+
_LOGGER.info("Discovery done")
143+
return # ignore timeouts on discover
144+
except Exception as ex:
145+
_LOGGER.warning("error while reading discover results: %s", ex)
146+
break
147147

148148
def send(
149149
self,
@@ -187,80 +187,80 @@ def send(
187187
Message.parse(m, token=self.token),
188188
)
189189

190-
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
191-
s.settimeout(self._timeout)
190+
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
191+
s.settimeout(self._timeout)
192192

193-
try:
194-
s.sendto(m, (self.ip, self.port))
195-
except OSError as ex:
196-
_LOGGER.error("failed to send msg: %s", ex)
197-
raise DeviceException from ex
198-
199-
try:
200-
data, addr = s.recvfrom(4096)
201-
m = Message.parse(data, token=self.token)
193+
try:
194+
s.sendto(m, (self.ip, self.port))
195+
except OSError as ex:
196+
_LOGGER.error("failed to send msg: %s", ex)
197+
raise DeviceException from ex
202198

203-
if self.debug > 1:
204-
_LOGGER.debug("recv from %s: %s", addr[0], m)
199+
try:
200+
data, addr = s.recvfrom(4096)
201+
m = Message.parse(data, token=self.token)
205202

206-
header = m.header.value
207-
payload = m.data.value
203+
if self.debug > 1:
204+
_LOGGER.debug("recv from %s: %s", addr[0], m)
208205

209-
self.__id = payload["id"]
210-
self._device_ts = header["ts"] # type: ignore # ts uses timeadapter
206+
header = m.header.value
207+
payload = m.data.value
211208

212-
_LOGGER.debug(
213-
"%s:%s (ts: %s, id: %s) << %s",
214-
self.ip,
215-
self.port,
216-
header["ts"],
217-
payload["id"],
218-
pf(payload),
219-
)
220-
if "error" in payload:
221-
self._handle_error(payload["error"])
209+
self.__id = payload["id"]
210+
self._device_ts = header["ts"] # type: ignore # ts uses timeadapter
222211

223-
try:
224-
return payload["result"]
225-
except KeyError:
226-
return payload
227-
except construct.core.ChecksumError as ex:
228-
raise InvalidTokenException(
229-
"Got checksum error which indicates use "
230-
"of an invalid token. "
231-
"Please check your token!"
232-
) from ex
233-
except OSError as ex:
234-
if retry_count > 0:
235212
_LOGGER.debug(
236-
"Retrying with incremented id, retries left: %s", retry_count
237-
)
238-
self.__id += 100
239-
self._discovered = False
240-
return self.send(
241-
command,
242-
parameters,
243-
retry_count - 1,
244-
extra_parameters=extra_parameters,
213+
"%s:%s (ts: %s, id: %s) << %s",
214+
self.ip,
215+
self.port,
216+
header["ts"],
217+
payload["id"],
218+
pf(payload),
245219
)
220+
if "error" in payload:
221+
self._handle_error(payload["error"])
222+
223+
try:
224+
return payload["result"]
225+
except KeyError:
226+
return payload
227+
except construct.core.ChecksumError as ex:
228+
raise InvalidTokenException(
229+
"Got checksum error which indicates use "
230+
"of an invalid token. "
231+
"Please check your token!"
232+
) from ex
233+
except OSError as ex:
234+
if retry_count > 0:
235+
_LOGGER.debug(
236+
"Retrying with incremented id, retries left: %s", retry_count
237+
)
238+
self.__id += 100
239+
self._discovered = False
240+
return self.send(
241+
command,
242+
parameters,
243+
retry_count - 1,
244+
extra_parameters=extra_parameters,
245+
)
246246

247-
_LOGGER.error("Got error when receiving: %s", ex)
248-
raise DeviceException("No response from the device") from ex
247+
_LOGGER.error("Got error when receiving: %s", ex)
248+
raise DeviceException("No response from the device") from ex
249249

250-
except RecoverableError as ex:
251-
if retry_count > 0:
252-
_LOGGER.debug(
253-
"Retrying to send failed command, retries left: %s", retry_count
254-
)
255-
return self.send(
256-
command,
257-
parameters,
258-
retry_count - 1,
259-
extra_parameters=extra_parameters,
260-
)
250+
except RecoverableError as ex:
251+
if retry_count > 0:
252+
_LOGGER.debug(
253+
"Retrying to send failed command, retries left: %s", retry_count
254+
)
255+
return self.send(
256+
command,
257+
parameters,
258+
retry_count - 1,
259+
extra_parameters=extra_parameters,
260+
)
261261

262-
_LOGGER.error("Got error when receiving: %s", ex)
263-
raise DeviceException("Unable to recover failed command") from ex
262+
_LOGGER.error("Got error when receiving: %s", ex)
263+
raise DeviceException("Unable to recover failed command") from ex
264264

265265
@property
266266
def _id(self) -> int:

‎miio/push_server/server.py‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -210,12 +210,13 @@ async def unsubscribe_event(self, device: Device, event_id: str):
210210

211211
async def _get_server_ip(self):
212212
"""Connect to the miio device to get server_ip using a one time use socket."""
213-
get_ip_socket = socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM)
214-
get_ip_socket.bind((self._address, SERVER_PORT))
215-
get_ip_socket.setblocking(False)
216-
await self._loop.sock_connect(get_ip_socket, (self._device_ip, SERVER_PORT))
217-
server_ip = get_ip_socket.getsockname()[0]
218-
get_ip_socket.close()
213+
with socket.socket(
214+
family=socket.AF_INET, type=socket.SOCK_DGRAM
215+
) as get_ip_socket:
216+
get_ip_socket.bind((self._address, SERVER_PORT))
217+
get_ip_socket.setblocking(False)
218+
await self._loop.sock_connect(get_ip_socket, (self._device_ip, SERVER_PORT))
219+
server_ip = get_ip_socket.getsockname()[0]
219220
_LOGGER.debug("Miio push server device ip=%s", server_ip)
220221
return server_ip
221222

0 commit comments

Comments
 (0)