goodwe 0.4.7__py3-none-any.whl → 0.4.9__py3-none-any.whl

goodwe/protocol.py

@@ -1,3 +1,4 @@
+"""Low level IP communication protocol implementation."""
 from __future__ import annotations
 
 import asyncio
@@ -6,7 +7,7 @@ import logging
 import platform
 import socket
 from asyncio.futures import Future
-from typing import Tuple, Optional, Callable
+from typing import Optional, Callable
 
 from .exceptions import MaxRetriesException, PartialResponseException, RequestFailedException, RequestRejectedException
 from .modbus import create_modbus_rtu_request, create_modbus_rtu_multi_request, create_modbus_tcp_request, \
@@ -37,7 +38,7 @@ class InverterProtocol:
         self._timer: asyncio.TimerHandle | None = None
         self.timeout: int = timeout
         self.retries: int = retries
-        self.keep_alive: bool = True
+        self.keep_alive: bool = False
         self.protocol: asyncio.Protocol | None = None
         self.response_future: Future | None = None
         self.command: ProtocolCommand | None = None
@@ -55,12 +56,29 @@ class InverterProtocol:
         """
         if self._lock and self._running_loop == asyncio.get_event_loop():
             return self._lock
-        else:
-            logger.debug("Creating lock instance for current event loop.")
-            self._lock = asyncio.Lock()
-            self._running_loop = asyncio.get_event_loop()
-            self._close_transport()
-            return self._lock
+        logger.debug("Creating lock instance for current event loop.")
+        self._lock = asyncio.Lock()
+        self._running_loop = asyncio.get_event_loop()
+        self._close_transport()
+        return self._lock
+
+    def _max_retries_reached(self) -> Future:
+        logger.debug("Max number of retries (%d) reached, request %s failed.", self.retries, self.command)
+        self._close_transport()
+        self.response_future = asyncio.get_running_loop().create_future()
+        self.response_future.set_exception(MaxRetriesException)
+        return self.response_future
+
+    def _close_transport(self) -> None:
+        if self._transport:
+            try:
+                self._transport.close()
+            except RuntimeError:
+                logger.debug("Failed to close transport.")
+            self._transport = None
+        # Cancel Future on connection lost
+        if self.response_future and not self.response_future.done():
+            self.response_future.cancel()
 
     async def close(self) -> None:
         """Close the underlying transport/connection."""
@@ -103,9 +121,11 @@ class UdpInverterProtocol(InverterProtocol, asyncio.DatagramProtocol):
 
     async def _connect(self) -> None:
         if not self._transport or self._transport.is_closing():
+            allow_broadcast = platform.system() == "Darwin" and self._host == "255.255.255.255"
             self._transport, self.protocol = await asyncio.get_running_loop().create_datagram_endpoint(
                 lambda: self,
                 remote_addr=(self._host, self._port),
+                allow_broadcast=allow_broadcast,
             )
 
     def connection_made(self, transport: asyncio.DatagramTransport) -> None:
@@ -120,7 +140,7 @@ class UdpInverterProtocol(InverterProtocol, asyncio.DatagramProtocol):
             logger.debug("Socket closed.")
         self._close_transport()
 
-    def datagram_received(self, data: bytes, addr: Tuple[str, int]) -> None:
+    def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None:
         """On datagram received"""
         if self._timer:
             self._timer.cancel()
@@ -133,15 +153,16 @@ class UdpInverterProtocol(InverterProtocol, asyncio.DatagramProtocol):
                 self._partial_missing = 0
             if self.command.validator(data):
                 logger.debug("Received: %s", data.hex())
+                self._retry = 0
                 self.response_future.set_result(data)
             else:
                 logger.debug("Received invalid response: %s", data.hex())
-                asyncio.get_running_loop().call_soon(self._retry_mechanism)
+                asyncio.get_running_loop().call_soon(self._timeout_mechanism)
         except PartialResponseException as ex:
             logger.debug("Received response fragment (%d of %d): %s", ex.length, ex.expected, data.hex())
             self._partial_data = data
             self._partial_missing = ex.expected - ex.length
-            self._timer = asyncio.get_running_loop().call_later(self.timeout, self._retry_mechanism)
+            self._timer = asyncio.get_running_loop().call_later(self.timeout, self._timeout_mechanism)
         except asyncio.InvalidStateError:
             logger.debug("Response already handled: %s", data.hex())
         except RequestRejectedException as ex:
@@ -158,13 +179,27 @@ class UdpInverterProtocol(InverterProtocol, asyncio.DatagramProtocol):
 
     async def send_request(self, command: ProtocolCommand) -> Future:
         """Send message via transport"""
-        async with self._ensure_lock():
+        await self._ensure_lock().acquire()
+        try:
             await self._connect()
             response_future = asyncio.get_running_loop().create_future()
-            self._retry = 0
             self._send_request(command, response_future)
             await response_future
             return response_future
+        except asyncio.CancelledError:
+            if self._retry < self.retries:
+                self._retry += 1
+                if self._lock and self._lock.locked():
+                    self._lock.release()
+                if not self.keep_alive:
+                    self._close_transport()
+                return await self.send_request(command)
+            return self._max_retries_reached()
+        finally:
+            if self._lock and self._lock.locked():
+                self._lock.release()
+            if not self.keep_alive:
+                self._close_transport()
 
     def _send_request(self, command: ProtocolCommand, response_future: Future) -> None:
         """Send message via transport"""
@@ -178,32 +213,19 @@ class UdpInverterProtocol(InverterProtocol, asyncio.DatagramProtocol):
         else:
             logger.debug("Sending: %s", self.command)
         self._transport.sendto(payload)
-        self._timer = asyncio.get_running_loop().call_later(self.timeout, self._retry_mechanism)
+        self._timer = asyncio.get_running_loop().call_later(self.timeout, self._timeout_mechanism)
 
-    def _retry_mechanism(self) -> None:
-        """Retry mechanism to prevent hanging transport"""
-        if self.response_future.done():
+    def _timeout_mechanism(self) -> None:
+        """Timeout mechanism to prevent hanging transport"""
+        if self.response_future and self.response_future.done():
             logger.debug("Response already received.")
-        elif self._retry < self.retries:
+            self._retry = 0
+        else:
             if self._timer:
                 logger.debug("Failed to receive response to %s in time (%ds).", self.command, self.timeout)
-            self._retry += 1
-            self._send_request(self.command, self.response_future)
-        else:
-            logger.debug("Max number of retries (%d) reached, request %s failed.", self.retries, self.command)
-            self.response_future.set_exception(MaxRetriesException)
-            self._close_transport()
-
-    def _close_transport(self) -> None:
-        if self._transport:
-            try:
-                self._transport.close()
-            except RuntimeError:
-                logger.debug("Failed to close transport.")
-            self._transport = None
-        # Cancel Future on connection close
-        if self.response_future and not self.response_future.done():
-            self.response_future.cancel()
+                self._timer = None
+            if self.response_future and not self.response_future.done():
+                self.response_future.cancel()
 
     async def close(self):
         self._close_transport()
@@ -250,7 +272,6 @@ class TcpInverterProtocol(InverterProtocol, asyncio.Protocol):
     def connection_made(self, transport: asyncio.DatagramTransport) -> None:
         """On connection made"""
         logger.debug("Connection opened.")
-        pass
 
     def eof_received(self) -> None:
         logger.debug("EOF received.")
@@ -319,8 +340,7 @@ class TcpInverterProtocol(InverterProtocol, asyncio.Protocol):
                     self._lock.release()
                 self._close_transport()
                 return await self.send_request(command)
-            else:
-                return self._max_retries_reached()
+            return self._max_retries_reached()
         except (ConnectionRefusedError, TimeoutError, OSError, asyncio.TimeoutError):
             if self._retry < self.retries:
                 logger.debug("Connection refused error.")
@@ -328,8 +348,7 @@ class TcpInverterProtocol(InverterProtocol, asyncio.Protocol):
                 if self._lock and self._lock.locked():
                     self._lock.release()
                 return await self.send_request(command)
-            else:
-                return self._max_retries_reached()
+            return self._max_retries_reached()
         finally:
             if self._lock and self._lock.locked():
                 self._lock.release()
@@ -358,24 +377,6 @@ class TcpInverterProtocol(InverterProtocol, asyncio.Protocol):
                 self._timer = None
             self._close_transport()
 
-    def _max_retries_reached(self) -> Future:
-        logger.debug("Max number of retries (%d) reached, request %s failed.", self.retries, self.command)
-        self._close_transport()
-        self.response_future = asyncio.get_running_loop().create_future()
-        self.response_future.set_exception(MaxRetriesException)
-        return self.response_future
-
-    def _close_transport(self) -> None:
-        if self._transport:
-            try:
-                self._transport.close()
-            except RuntimeError:
-                logger.debug("Failed to close transport.")
-            self._transport = None
-        # Cancel Future on connection lost
-        if self.response_future and not self.response_future.done():
-            self.response_future.cancel()
-
     async def close(self):
         await self._ensure_lock().acquire()
         try:
@@ -399,8 +400,7 @@ class ProtocolResponse:
     def response_data(self) -> bytes:
         if self.command is not None:
             return self.command.trim_response(self.raw_data)
-        else:
-            return self.raw_data
+        return self.raw_data
 
     def seek(self, address: int) -> None:
         if self.command is not None:
@@ -454,10 +454,9 @@ class ProtocolCommand:
             result = response_future.result()
             if result is not None:
                 return ProtocolResponse(result, self)
-            else:
-                raise RequestFailedException(
-                    "No response received to '" + self.request.hex() + "' request."
-                )
+            raise RequestFailedException(
+                "No response received to '" + self.request.hex() + "' request."
+            )
         except (asyncio.CancelledError, ConnectionRefusedError):
             raise RequestFailedException(
                 "No valid response received to '" + self.request.hex() + "' request."
@@ -540,12 +539,11 @@ class Aa55ProtocolCommand(ProtocolCommand):
         if self.request[4] == 1:
             if self.request[5] == 2:
                 return f'READ device info ({self.request.hex()})'
-            elif self.request[5] == 6:
+            if self.request[5] == 6:
                 return f'READ runtime data ({self.request.hex()})'
-            elif self.request[5] == 9:
+            if self.request[5] == 9:
                 return f'READ settings ({self.request.hex()})'
-        else:
-            return self.request.hex()
+        return self.request.hex()
 
 
 class Aa55ReadCommand(Aa55ProtocolCommand):
@@ -554,13 +552,12 @@ class Aa55ReadCommand(Aa55ProtocolCommand):
     """
 
     def __init__(self, offset: int, count: int):
-        super().__init__("011A03" + "{:04x}".format(offset) + "{:02x}".format(count), "019A", offset, count)
+        super().__init__(f"011A03{offset:04x}{count:02x}", "019A", offset, count)
 
     def __repr__(self):
         if self.value > 1:
             return f'READ {self.value} registers from {self.first_address} ({self.request.hex()})'
-        else:
-            return f'READ register {self.first_address} ({self.request.hex()})'
+        return f'READ register {self.first_address} ({self.request.hex()})'
 
 
 class Aa55WriteCommand(Aa55ProtocolCommand):
@@ -569,7 +566,7 @@ class Aa55WriteCommand(Aa55ProtocolCommand):
     """
 
     def __init__(self, register: int, value: int):
-        super().__init__("023905" + "{:04x}".format(register) + "01" + "{:04x}".format(value), "02B9", register, value)
+        super().__init__(f"023905{register:04x}01{value:04x}", "02B9", register, value)
 
     def __repr__(self):
         return f'WRITE {self.value} to register {self.first_address} ({self.request.hex()})'
@@ -581,7 +578,7 @@ class Aa55WriteMultiCommand(Aa55ProtocolCommand):
     """
 
     def __init__(self, offset: int, values: bytes):
-        super().__init__("02390B" + "{:04x}".format(offset) + "{:02x}".format(len(values)) + values.hex(),
+        super().__init__(f"02390B{offset:04x}{len(values):02x}{values.hex()}",
                          "02B9", offset, len(values) // 2)
 
 
@@ -637,8 +634,7 @@ class ModbusRtuReadCommand(ModbusRtuProtocolCommand):
     def __repr__(self):
         if self.value > 1:
             return f'READ {self.value} registers from {self.first_address} ({self.request.hex()})'
-        else:
-            return f'READ register {self.first_address} ({self.request.hex()})'
+        return f'READ register {self.first_address} ({self.request.hex()})'
 
 
 class ModbusRtuWriteCommand(ModbusRtuProtocolCommand):
@@ -707,8 +703,7 @@ class ModbusTcpReadCommand(ModbusTcpProtocolCommand):
     def __repr__(self):
         if self.value > 1:
             return f'READ {self.value} registers from {self.first_address} ({self.request.hex()})'
-        else:
-            return f'READ register {self.first_address} ({self.request.hex()})'
+        return f'READ register {self.first_address} ({self.request.hex()})'
 
 
 class ModbusTcpWriteCommand(ModbusTcpProtocolCommand):