Fix conditional substructures

suitcase/fields.py

@@ -612,11 +612,14 @@ def pack(self, stream):
         if self.condition(self._parent):
             self.field.pack(stream)
 
-    def unpack(self, data):
+    def unpack(self, data, **kwargs):
         if self.condition(self._parent):
-            self.field.unpack(data)
+            return self.field.unpack(data, **kwargs)
 
     def getval(self):
+        if not self.condition(self._parent):
+            return None
+
         return self.field.getval()
 
     def setval(self, value):

suitcase/structure.py

@@ -8,7 +8,8 @@
 import six
 from suitcase.exceptions import SuitcaseException, \
     SuitcasePackException, SuitcaseParseError
-from suitcase.fields import FieldPlaceholder, CRCField, SubstructureField
+from suitcase.fields import FieldPlaceholder, CRCField, SubstructureField, \
+    ConditionalField
 from six import BytesIO
 
 
@@ -85,6 +86,11 @@ def unpack_stream(self, stream):
         logic present for dealing with checksum fields.
 
         """
+        def is_substructure(field):
+            return isinstance(field, SubstructureField) or \
+                    isinstance(field, ConditionalField) and \
+                    isinstance(field.field, SubstructureField)
+
         crc_fields = []
         greedy_field = None
         # go through the fields from first to last.  If we hit a greedy
@@ -93,11 +99,15 @@ def unpack_stream(self, stream):
             if isinstance(field, CRCField):
                 crc_fields.append((field, stream.tell()))
             length = field.bytes_required
-            if isinstance(field, SubstructureField):
+            if is_substructure(field):
                 remaining_data = stream.getvalue()[stream.tell():]
                 returned_stream = field.unpack(remaining_data, trailing=True)
+                if returned_stream is not None:
+                    consumed = returned_stream.tell()
+                else:
+                    consumed = 0
                 # We need to fast forward by as much as was consumed by the structure
-                stream.seek(stream.tell() + returned_stream.tell())
+                stream.seek(stream.tell() + consumed)
                 continue
             elif length is None:
                 greedy_field = field
@@ -135,7 +145,7 @@ def unpack_stream(self, stream):
                     raise SuitcaseParseError("While attempting to parse field "
                                              "%r we tried to read %s bytes but "
                                              "we were only able to read %s." %
-                                             (name, length, len(data)))
+                                             (_name, length, len(data)))
                 try:
                     field.unpack(data)
                 except SuitcaseException:

suitcase/test/test_fields.py

@@ -668,6 +668,14 @@ class ConditionalDispatch(Structure):
             condition=lambda m: m.f2 == 255)
 
 
+class ConditionalSubstructures(Structure):
+    f1 = UBInt8()
+    f2 = ConditionalField(SubstructureField(BasicMessage),
+                          condition=lambda m: m.f1 < 10)
+    f3 = ConditionalField(SubstructureField(BasicMessage),
+                          condition=lambda m: m.f1 < 20)
+
+
 class TestConditionalField(unittest.TestCase):
     def test_conditional_pack(self):
         m1 = Conditional()
@@ -732,6 +740,62 @@ def test_conditional_dispatch(self):
         self.assertEqual(m2.f3.b1, 0x11)
         self.assertEqual(m2.f3.b2, 0x22)
 
+    def test_conditional_substructure_pack(self):
+        m = ConditionalSubstructures()
+        m.f1 = 9
+
+        f2 = BasicMessage()
+        f2.b1 = 0x55
+        f2.b2 = 0x66
+        m.f2 = f2
+
+        f3 = BasicMessage()
+        f3.b1 = 0x77
+        f3.b2 = 0x88
+        m.f3 = f3
+
+        self.assertEqual(m.pack(), b"\x09\x55\x66\x77\x88")
+
+        # Remove f2
+        m.f1 = 10
+        m.f2 = None
+        self.assertEqual(m.pack(), b"\x0a\x77\x88")
+
+        # Remove f3
+        m.f1 = 20
+        m.f3 = None
+        self.assertEqual(m.pack(), b"\x14")
+
+    def test_conditional_substructure_unpack(self):
+        # Neither substructure is present
+        m1 = ConditionalSubstructures.from_data(b"\x20")
+        self.assertEqual(m1.f1, 0x20)
+        self.assertEqual(m1.f2, None)
+        self.assertEqual(m1.f3, None)
+
+        # Include f3 but not f2
+        m2 = ConditionalSubstructures()
+        m2.unpack(b"\x10" +
+                  b"\x11\x22")
+        self.assertEqual(m2.f1, 0x10)
+        self.assertEqual(m2.f2, None)
+        self.assertIsInstance(m2.f3, BasicMessage)
+        self.assertEqual(m2.f3.b1, 0x11)
+        self.assertEqual(m2.f3.b2, 0x22)
+
+        # Include both substructures
+        m3 = ConditionalSubstructures()
+        m3.unpack(b"\x00" +
+                  b"\x11\x22" +
+                  b"\x33\x44")
+        self.assertEqual(m3.f1, 0x00)
+        self.assertIsInstance(m3.f2, BasicMessage)
+        self.assertEqual(m3.f2.b1, 0x11)
+        self.assertEqual(m3.f2.b2, 0x22)
+        self.assertIsInstance(m3.f3, BasicMessage)
+        self.assertEqual(m3.f3.b1, 0x33)
+        self.assertEqual(m3.f3.b2, 0x44)
+
 
 class TestStructure(unittest.TestCase):
     def test_unpack_fewer_bytes_than_required(self):