Be more careful with type checking for framer

pymodbus-dev · Dec 4, 2024 · 39d8096 · 39d8096
1 parent 7c964f4
commit 39d8096
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/pymodbus/client/base.py b/pymodbus/client/base.py
@@ -35,9 +35,11 @@ def __init__(
         """
         ModbusClientMixin.__init__(self)  # type: ignore[arg-type]
         self.comm_params = comm_params
+        if isinstance(framer, FramerType):
+            framer = FRAMER_NAME_TO_CLASS[framer]
         self.ctx = TransactionManager(
             comm_params,
-            FRAMER_NAME_TO_CLASS.get(framer, framer)(DecodePDU(False)),
+            framer(DecodePDU(False)),
             retries,
             False,
             trace_packet,
@@ -133,7 +135,9 @@ def __init__(
         self.slaves: list[int] = []
 
         # Common variables.
-        self.framer: FramerBase = FRAMER_NAME_TO_CLASS.get(framer, framer)(DecodePDU(False))
+        if isinstance(framer, FramerType):
+            framer = FRAMER_NAME_TO_CLASS[framer]
+        self.framer: FramerBase = framer(DecodePDU(False))
         self.transaction = TransactionManager(
             self.comm_params,
             self.framer,

diff --git a/pymodbus/server/async_io.py b/pymodbus/server/async_io.py
@@ -10,7 +10,7 @@
 from pymodbus.datastore import ModbusServerContext
 from pymodbus.device import ModbusControlBlock, ModbusDeviceIdentification
 from pymodbus.exceptions import NoSuchSlaveException
-from pymodbus.framer import FRAMER_NAME_TO_CLASS, FramerType, FramerBase
+from pymodbus.framer import FRAMER_NAME_TO_CLASS, FramerBase, FramerType
 from pymodbus.logging import Log
 from pymodbus.pdu import DecodePDU, ModbusPDU
 from pymodbus.pdu.pdu import ExceptionResponse
@@ -224,7 +224,9 @@ def __init__(
         if isinstance(identity, ModbusDeviceIdentification):
             self.control.Identity.update(identity)
         # Support mapping of FramerType to a Framer class, or a Framer class
-        self.framer = FRAMER_NAME_TO_CLASS.get(framer, framer)
+        if isinstance(framer, FramerType):
+            framer = FRAMER_NAME_TO_CLASS[framer]
+        self.framer = framer
         self.serving: asyncio.Future = asyncio.Future()
 
     def callback_new_connection(self):
@@ -273,7 +275,7 @@ def __init__(
         self,
         context: ModbusServerContext,
         *,
-        framer=FramerType.SOCKET,
+        framer: FramerType | type[FramerBase] = FramerType.SOCKET,
         identity: ModbusDeviceIdentification | None = None,
         address: tuple[str, int] = ("", 502),
         ignore_missing_slaves: bool = False,
@@ -336,7 +338,7 @@ def __init__(  # pylint: disable=too-many-arguments
         self,
         context: ModbusServerContext,
         *,
-        framer=FramerType.TLS,
+        framer: FramerType | type[FramerBase] = FramerType.TLS,
         identity: ModbusDeviceIdentification | None = None,
         address: tuple[str, int] = ("", 502),
         sslctx=None,
@@ -403,7 +405,7 @@ def __init__(
         self,
         context: ModbusServerContext,
         *,
-        framer=FramerType.SOCKET,
+        framer: FramerType | type[FramerBase] = FramerType.SOCKET,
         identity: ModbusDeviceIdentification | None = None,
         address: tuple[str, int] = ("", 502),
         ignore_missing_slaves: bool = False,
@@ -463,7 +465,7 @@ def __init__(
         self,
         context: ModbusServerContext,
         *,
-        framer: FramerType = FramerType.RTU,
+        framer: FramerType | type[FramerBase] = FramerType.RTU,
         ignore_missing_slaves: bool = False,
         identity: ModbusDeviceIdentification | None = None,
         broadcast_enable: bool = False,