From 936a6b30f41c8e016a6ce4678279fb781f6a910f Mon Sep 17 00:00:00 2001
From: Mike Wadsten <mikewadsten@gmail.com>
Date: Thu, 19 Nov 2015 20:32:50 -0600
Subject: [PATCH] fields: Make ConditionalField.getval return None when
 condition is not met

---
 suitcase/fields.py           |  3 ++
 suitcase/test/test_fields.py | 64 ++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/suitcase/fields.py b/suitcase/fields.py
index e0e8e2f..89556e1 100644
--- a/suitcase/fields.py
+++ b/suitcase/fields.py
@@ -617,6 +617,9 @@ def unpack(self, data, **kwargs):
             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):
diff --git a/suitcase/test/test_fields.py b/suitcase/test/test_fields.py
index 6482b31..da8d19a 100644
--- a/suitcase/test/test_fields.py
+++ b/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):