Skip to content

Commit

Permalink
Merge 936a6b3 into acdfb85
Browse files Browse the repository at this point in the history
  • Loading branch information
mikewadsten committed Nov 20, 2015
2 parents acdfb85 + 936a6b3 commit 6e41b91
Show file tree
Hide file tree
Showing 3 changed files with 83 additions and 6 deletions.
7 changes: 5 additions & 2 deletions suitcase/fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
18 changes: 14 additions & 4 deletions suitcase/structure.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand Down
64 changes: 64 additions & 0 deletions suitcase/test/test_fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit 6e41b91

Please sign in to comment.