Merge pull request #35 from rcfox/fieldarray-num-elements

Accept a num_elements_provider in FieldArray to specify a size in elements instead of bytes.

Author: posborne

suitcase/fields.py

@@ -845,6 +845,8 @@ class FieldArray(BaseField):
         the entire contents of the array.
     :param length_provider: The field providing a length value binding this
         message (if any).  Set this field to None to leave unconstrained.
+    :param num_elements_provider: The field providing the number of elements
+        in the array. Set this field to None to leave unconstrained.
 
     For example, one can imagine a message listing the zipcodes covered by
     a telephone area code. Depending on the population density, the number
@@ -863,7 +865,7 @@ class TelephoneZipcodes(Structure):
 
     """
 
-    def __init__(self, substructure, length_provider=None, **kwargs):
+    def __init__(self, substructure, length_provider=None, num_elements_provider=None, **kwargs):
         BaseField.__init__(self, **kwargs)
         self.substructure = substructure
         self._value = list()
@@ -872,6 +874,12 @@ def __init__(self, substructure, length_provider=None, **kwargs):
             self.length_provider.associate_length_consumer(self)
         else:
             self.length_provider = None
+        if isinstance(num_elements_provider, FieldPlaceholder):
+            self.num_elements_provider = self._ph2f(num_elements_provider)
+            self.num_elements_provider.length_value_provider = lambda: len(self._value)
+        else:
+            self.num_elements_provider = None
+
 
     @property
     def bytes_required(self):
@@ -880,15 +888,22 @@ def bytes_required(self):
         else:
             return self.length_provider.get_adjusted_length()
 
+    @property
+    def num_elements(self):
+        if self.num_elements_provider is None:
+            return None
+        else:
+            return self.num_elements_provider.get_adjusted_length()
+
     def pack(self, stream):
         for structure in self._value:
             stream.write(structure.pack())
 
     def unpack(self, data, **kwargs):
         length = self.bytes_required
-        if length == 0 or (data == b"" and length is None):
+        if length == 0 or (data == b"" and length is None) or self.num_elements == 0:
             # Array is empty.
-            return
+            return data
 
         kwargs['trailing'] = True
         while True:
@@ -897,6 +912,14 @@ def unpack(self, data, **kwargs):
             self._value.append(structure)
             if data == b"":
                 break
+            if len(self._value) == self.num_elements:
+                break
+
+        if self.num_elements is not None and len(self._value) != self.num_elements:
+            raise SuitcaseParseError("Expected %s elements but received %d." %
+                                     (self.num_elements, len(self._value)))
+
+        return data
 
 
 class BaseFixedByteSequence(BaseField):

suitcase/structure.py

@@ -4,11 +4,11 @@
 #
 # Copyright (c) 2015 Digi International Inc. All Rights Reserved.
 import sys
-
+import os
 import six
 from suitcase.exceptions import SuitcaseException, \
     SuitcasePackException, SuitcaseParseError
-from suitcase.fields import FieldPlaceholder, CRCField, SubstructureField, \
+from suitcase.fields import FieldArray, FieldPlaceholder, CRCField, SubstructureField, \
     ConditionalField
 from six import BytesIO
 
@@ -107,8 +107,12 @@ def unpack_stream(self, stream):
                 stream.seek(stream.tell() + consumed)
                 continue
             elif length is None:
-                greedy_field = field
-                break
+                if isinstance(field, FieldArray) and field.num_elements is not None:
+                    # Read the data greedily now, and we'll backtrack after enough elements have been read.
+                    data = stream.read()
+                else:
+                    greedy_field = field
+                    break
             else:
                 data = stream.read(length)
                 if len(data) != length:
@@ -118,7 +122,8 @@ def unpack_stream(self, stream):
                                              (name, length, len(data)))
 
             try:
-                field.unpack(data)
+                unused_data = field.unpack(data)
+                stream.seek(-len(unused_data or ""), os.SEEK_CUR)
             except SuitcaseException:
                 raise  # just re-raise these
             except Exception:

suitcase/test/test_fields.py

@@ -1043,6 +1043,12 @@ class BasicMessageArrayAfter(Structure):
     after = UBInt8()
 
 
+# Test a FieldArray that takes a number of elements instead of a size in bytes.
+class BasicMessageArrayNumElements(Structure):
+    count = LengthField(UBInt8())
+    array = FieldArray(BasicMessage, num_elements_provider=count)
+
+
 class TestFieldArray(unittest.TestCase):
     def test_pack_valid(self):
         m = BasicMessageArray()
@@ -1112,6 +1118,30 @@ def test_unpack_after_valid(self):
         self.assertEqual(m2.array[0].b2, 0x22)
         self.assertEqual(m2.after, 0)
 
+    def test_pack_num_elements_valid(self):
+        m = BasicMessageArrayNumElements()
+        m.array = []
+        self.assertEqual(m.pack(), b"\x00")
+
+        messages = [BasicMessage(), BasicMessage(), BasicMessage()]
+        for i, message in enumerate(messages):
+            message.b1 = 0x10 + i
+            message.b2 = 0x20 + i
+            m.array.append(message)
+
+        self.assertEqual(m.pack(), b"\x03\x10\x20\x11\x21\x12\x22")
+
+    def test_unpack_num_elements_valid(self):
+        m = BasicMessageArrayNumElements.from_data(b"\x02\x22\x33\x44\x55")
+
+        self.assertEqual(m.count, 2)
+        self.assertEqual(len(m.array), 2)
+        self.assertIsInstance(m.array[0], BasicMessage)
+        self.assertEqual(m.array[0].b1, 0x22)
+        self.assertEqual(m.array[0].b2, 0x33)
+        self.assertEqual(m.array[1].b1, 0x44)
+        self.assertEqual(m.array[1].b2, 0x55)
+
 
 # Test SubstructureField
 class PascalString16(Structure):