Fix mypy errors in strict mode (#28)

Author: nikteliy

.pre-commit-config.yaml

@@ -16,10 +16,12 @@ repos:
       - id: check-docstring-first
       - id: detect-private-key
 
-  # - repo: https://github.com/pre-commit/mirrors-mypy
-  #   rev: 'v1.10.1'
-  #   hooks:
-  #     - id: mypy
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: 'v1.10.1'
+    hooks:
+      - id: mypy
+        additional_dependencies:
+          - pytest
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: 'v0.5.1'

src/python_s7comm/async_client.py

@@ -1,5 +1,4 @@
 import datetime
-from typing import List, Optional, Tuple
 
 from .s7comm import AsyncS7Comm, enums
 from .s7comm.packets.variable_address import VariableAddress
@@ -50,31 +49,32 @@ async def get_cpu_state(self) -> enums.CPUStatus:
     async def read_area(self, address: str) -> bytes:
         return await self.s7comm.read_area(VariableAddress.from_string(address))
 
-    async def write_area(self, address: str, data: bytes) -> int:
-        return await self.s7comm.write_area(address=VariableAddress.from_string(address), data=data)
+    async def write_area(self, address: str, data: bytes) -> None:
+        await self.s7comm.write_area(address=VariableAddress.from_string(address), data=data)
 
-    async def get_order_code(self) -> Optional[str]:
+    async def get_order_code(self) -> str | None:
         response = await self.s7comm.read_szl(szl_id=0x0011, szl_index=0x0000)
-        szl_data = SZLResponseData.parse(response.data)
+        szl_data = SZLResponseData.parse(response.data.data)
         for date_tree in szl_data.szl_data_tree_list:
             module_identification = ModuleIdentificationDataTree.parse(date_tree)
             if module_identification.index == ModuleIdentificationIndex.MODULE_IDENTIFICATION:
                 return module_identification.order_number
         return None
 
     async def read_szl(self, szl_id: int, szl_index: int = 0x0000) -> SZLResponseData:
-        return await self.s7comm.read_szl(szl_id=szl_id, szl_index=szl_index)
+        response_data = await self.s7comm.read_szl(szl_id=szl_id, szl_index=szl_index)
+        return SZLResponseData.parse(response_data.data.data)
 
     async def read_szl_list(self) -> SZLResponseData:
         response_data = await self.s7comm.read_szl(szl_id=0x0000, szl_index=0x0000)
         return SZLResponseData.parse(response_data.data.data)
 
-    async def read_multi_vars(self, items: List[str]) -> List[bytes]:
+    async def read_multi_vars(self, items: list[str]) -> list[bytes]:
         vars_ = [VariableAddress.from_string(item) for item in items]
         response = await self.s7comm.read_multi_vars(items=vars_)
         return response.values()
 
-    async def write_multi_vars(self, items: List[Tuple[str, bytes]]) -> bool:
+    async def write_multi_vars(self, items: list[tuple[str, bytes]]) -> bool:
         vars_ = [(VariableAddress.from_string(address), data) for address, data in items]
         response = await self.s7comm.write_multi_vars(items=vars_)
         return response.check_result()
@@ -85,7 +85,7 @@ async def set_plc_system_datetime(self) -> int:
     async def delete(self, block_type: str, block_num: int) -> int:
         raise NotImplementedError
 
-    async def full_upload(self, _type: str, block_num: int) -> Tuple[bytearray, int]:
+    async def full_upload(self, _type: str, block_num: int) -> tuple[bytearray, int]:
         raise NotImplementedError
 
     async def upload(self, block_num: int) -> bytearray:

src/python_s7comm/s7comm/__init__.py

@@ -1,5 +1,6 @@
 from .async_client import AsyncS7Comm
 from .client import S7Comm
+from .exceptions import PacketLostError, StalePacketError
 
 
-__all__ = ["S7Comm", "AsyncS7Comm"]
+__all__ = ["S7Comm", "AsyncS7Comm", "PacketLostError", "StalePacketError"]

src/python_s7comm/s7comm/async_client.py

@@ -4,14 +4,15 @@
 from typing import cast
 
 from .base import BaseS7Comm
-from .enums import HeaderError, ItemReturnCode, MessageType, SubfunctionCode, UserdataFunction
+from .enums import HeaderError, ItemReturnCode, MessageType, SubfunctionCode, UserdataFunction, UserdataLastPDU
 from .error_messages import ERROR_MESSAGES
 from .exceptions import PacketLostError, StalePacketError
 from .packets import (
     RequestPLCStop,
     S7AckDataHeader,
     S7Packet,
     SetupCommunicationRequest,
+    SZLResponseData,
     UserDataContinuationRequest,
     UserDataRequest,
     UserDataResponse,
@@ -38,6 +39,7 @@ def __init__(
         transport: AsyncTransport | None = None,
     ):
         super().__init__(pdu_length=pdu_length)
+        self.transport: AsyncTransport
         if transport is None:
             self.transport = AsyncCOTP(tpdu_size=tpdu_size, source_tsap=source_tsap, dest_tsap=dest_tsap)
         else:
@@ -132,7 +134,7 @@ async def read_area(self, address: VariableAddress) -> bytes:
                     result += response_item.data
         return result
 
-    async def write_area(self, address: VariableAddress, data: bytes) -> int:
+    async def write_area(self, address: VariableAddress, data: bytes) -> WriteVariableResponse:
         """
         Writes data to PLC memory area with automatic splitting into multiple requests
         if data exceeds PDU size.
@@ -177,7 +179,7 @@ async def write_multi_vars(self, items: list[tuple[VariableAddress, bytes]]) ->
             raise ValueError("Invalid response class")
         return response
 
-    async def read_szl(self, szl_id: int, szl_index: int) -> S7Packet:
+    async def read_szl(self, szl_id: int, szl_index: int) -> UserDataResponse:
         data = struct.pack("!HH", szl_id, szl_index)
         request = UserDataRequest.create(
             function_group=UserdataFunction.CPU_FUNCTION,
@@ -186,6 +188,8 @@ async def read_szl(self, szl_id: int, szl_index: int) -> S7Packet:
             data=data,
         )
         response = await self.send(request=request)
+        if not isinstance(response, UserDataResponse):
+            raise ValueError("Invalid response class")
         return response
 
     async def plc_stop(self) -> S7Packet:

src/python_s7comm/s7comm/base.py

@@ -47,6 +47,7 @@ def _create_packet(request: S7Packet, pdu_reference: int) -> bytes:
 
     def _validate_pdu_reference(self, response: S7Packet) -> None:
         """Raises if PDU reference is invalid."""
+        assert response.header is not None
         if response.header.pdu_reference > self._pdu_reference:
             raise PacketLostError(f"Expected {self._pdu_reference}, got {response.header.pdu_reference}")
         elif response.header.pdu_reference < self._pdu_reference:

src/python_s7comm/s7comm/client.py

@@ -37,6 +37,7 @@ def __init__(
         transport: Transport | None = None,
     ) -> None:
         super().__init__(pdu_length=pdu_length)
+        self.transport: Transport
         if transport is None:
             self.transport = COTP(tpdu_size=tpdu_size, source_tsap=source_tsap, dest_tsap=dest_tsap)
         else:

src/python_s7comm/s7comm/packets/__init__.py

@@ -1,19 +1,23 @@
+from .data_item import DataItem
 from .headers import S7AckDataHeader, S7Header
 from .packet import S7Packet
 from .plc_command import RequestPLCStop
 from .rw_variable import VariableReadRequest, VariableWriteRequest
 from .setup_communication import SetupCommunicationParameter, SetupCommunicationRequest
+from .szl import SZLResponseData
 from .user_data import UserDataContinuationRequest, UserDataRequest, UserDataResponse
 from .variable_address import VariableAddress
 
 
 __all__ = [
+    "DataItem",
     "S7AckDataHeader",
     "S7Parameter",
     "RequestPLCStop",
     "SetupCommunicationRequest",
     "SetupCommunicationParameter",
     "S7Packet",
+    "SZLResponseData",
     "UserDataContinuationRequest",
     "UserDataRequest",
     "UserDataResponse",

src/python_s7comm/s7comm/packets/error.py

@@ -16,3 +16,7 @@ def serialize_parameter(self) -> bytes:
 
     def serialize_data(self) -> bytes:
         return b""
+
+    @classmethod
+    def parse(cls, packet: bytes) -> "S7Error":
+        raise NotImplementedError("S7Error is created from header, not parsed from bytes")

src/python_s7comm/s7comm/packets/packet.py

@@ -1,4 +1,4 @@
-from typing import Any, ClassVar, Protocol
+from typing import Any, ClassVar, Protocol, Self
 
 from .headers import S7AckDataHeader, S7Header
 
@@ -7,16 +7,23 @@ class S7Data(Protocol):
     def serialize(self) -> bytes: ...
 
 
+class S7Parameter(Protocol):
+    def serialize(self) -> bytes: ...
+
+
 class S7Packet(Protocol):
     MESSAGE_TYPE: ClassVar
     header: S7Header | None
-    parameter: Any  # Optional[S7Parameter]
-    data: Any  # Optional[S7Data]  # Union[DataItem, List[DataItem], List[enums.ItemReturnCode], None]
+    parameter: Any  # S7Parameter | None
+    data: Any  # S7Data | list[S7Data] | None
 
     def serialize_parameter(self) -> bytes: ...
 
     def serialize_data(self) -> bytes: ...
 
+    @classmethod
+    def parse(cls, packet: bytes) -> "S7Packet": ...
+
 
 class S7Response(S7Packet):
     def create_packet(self, pdu_reference: int) -> bytes:

src/python_s7comm/s7comm/packets/parser.py

@@ -14,13 +14,13 @@
 from .user_data import UserDataResponse
 
 
-response_code_to_packet: dict[FunctionCode, S7Packet] = {
+response_code_to_packet: dict[FunctionCode, type[S7Packet]] = {
     FunctionCode.SetupCommunication: SetupCommunicationRequest,
     FunctionCode.ReadVariable: ReadVariableResponse,
     FunctionCode.WriteVariable: WriteVariableResponse,
 }
 
-request_code_to_packet: dict[FunctionCode, S7Packet] = {
+request_code_to_packet: dict[FunctionCode, type[S7Packet]] = {
     FunctionCode.SetupCommunication: SetupCommunicationRequest,
     FunctionCode.ReadVariable: VariableReadRequest,
     FunctionCode.WriteVariable: VariableWriteRequest,
@@ -47,12 +47,12 @@ def parse(packet: bytes) -> S7Packet:
 
             # parse parameter
             parameter = packet[header.LENGTH :]
-            response = response_code_to_packet.get(function_code).parse(parameter)
+            response = response_code_to_packet[function_code].parse(parameter)
         elif message_type == MessageType.JobRequest:
             header = S7Header.parse(packet)
             function_code = struct.unpack_from("!B", packet, header.LENGTH)[0]
             parameter = packet[header.LENGTH :]
-            response = request_code_to_packet.get(function_code).parse(parameter)
+            response = request_code_to_packet[function_code].parse(parameter)
         elif message_type == MessageType.Userdata:
             header = S7Header.parse(packet)
             response = UserDataResponse.parse(packet[header.LENGTH :])