From c5a9eaeab363d234b9343794582d37e898a81c96 Mon Sep 17 00:00:00 2001
From: Cougar <cougar@random.ee>
Date: Fri, 15 Aug 2014 23:29:08 +0300
Subject: [PATCH] Implement getResponseSize()

getResponseSize() can be used to know in advance how many
bytes needs to be read after request. Using this information it
is not needed to wait for serial port timeout during receive.
---
 pymodbus/pdu.py                    |  7 +++++++
 pymodbus/register_read_message.py  |  7 +++++++
 pymodbus/register_write_message.py |  7 +++++++
 pymodbus/transaction.py            | 27 ++++++++++++++++++++++-----
 4 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/pymodbus/pdu.py b/pymodbus/pdu.py
index 596b4cc99..161173300 100644
--- a/pymodbus/pdu.py
+++ b/pymodbus/pdu.py
@@ -106,6 +106,13 @@ def doException(self, exception):
                 (self.function_code, exception))
         return ExceptionResponse(self.function_code, exception)
 
+    def getResponseSize(self):
+        ''' Returns expected packet size of response for this request
+
+        :raises: A not implemented exception
+        '''
+        raise NotImplementedException()
+
 
 class ModbusResponse(ModbusPDU):
     ''' Base class for a modbus response PDU
diff --git a/pymodbus/register_read_message.py b/pymodbus/register_read_message.py
index 06d82925f..76185beef 100644
--- a/pymodbus/register_read_message.py
+++ b/pymodbus/register_read_message.py
@@ -38,6 +38,13 @@ def decode(self, data):
         '''
         self.address, self.count = struct.unpack('>HH', data)
 
+    def getResponseSize(self):
+        ''' Returns expected packet size of response for this request
+
+        :returns: The expected packet size
+        '''
+        return 1 + 2 * self.count
+
     def __str__(self):
         ''' Returns a string representation of the instance
 
diff --git a/pymodbus/register_write_message.py b/pymodbus/register_write_message.py
index a1dcf91f7..160ae88ce 100644
--- a/pymodbus/register_write_message.py
+++ b/pymodbus/register_write_message.py
@@ -61,6 +61,13 @@ def execute(self, context):
         values = context.getValues(self.function_code, self.address, 1)
         return WriteSingleRegisterResponse(self.address, values[0])
 
+    def getResponseSize(self):
+        ''' Returns expected packet size of response for this request
+
+        :returns: The expected packet size
+        '''
+        return 2 + 2
+
     def __str__(self):
         ''' Returns a string representation of the instance
 
diff --git a/pymodbus/transaction.py b/pymodbus/transaction.py
index 5589e58da..78756f1c7 100644
--- a/pymodbus/transaction.py
+++ b/pymodbus/transaction.py
@@ -6,7 +6,7 @@
 import socket
 from binascii import b2a_hex, a2b_hex
 
-from pymodbus.exceptions import ModbusIOException
+from pymodbus.exceptions import ModbusIOException, NotImplementedException
 from pymodbus.constants  import Defaults
 from pymodbus.interfaces import IModbusFramer
 from pymodbus.utilities  import checkCRC, computeCRC
@@ -62,10 +62,19 @@ def execute(self, request):
             try:
                 self.client.connect()
                 self.client._send(self.client.framer.buildPacket(request))
-                # I need to fix this to read the header and the result size,
-                # as this may not read the full result set, but right now
-                # it should be fine...
-                result = self.client._recv(1024)
+                try:
+                    recvsize = self.client.framer.getResponseSize(request)
+                except NotImplementedException:
+                    recvsize = 0;
+
+                # TODO: read first only as much as needed to check Modbus
+                # exception and then all remaining bytes. Otherwise
+                # exception always takes _timeout_ seconds.
+                if recvsize:
+                    result = self.client._recv(recvsize)
+                else:
+                    result = self.client._recv(1024)
+
                 if not result and self.retry_on_empty:
                     retries -= 1
                     continue
@@ -553,6 +562,14 @@ def buildPacket(self, message):
         packet += struct.pack(">H", computeCRC(packet))
         return packet
 
+    def getResponseSize(self, message):
+        ''' Returns expected packet size of response for request
+
+        :param message: Request message
+        :returns: The expected packet size
+        '''
+        return 1 + 1 + message.getResponseSize() + 2
+
 
 #---------------------------------------------------------------------------#
 # Modbus ASCII Message