Support more than one nested struct.

Bug: 34812052
Test: new unit test passes
Change-Id: I7c3f4d73be01e5000e8b870e9178d6408993ae0d

2 files changed, 54 insertions, 2 deletions

diff --git a/net/test/cstruct.py b/net/test/cstruct.py
index 91cd72e..8fe916a 100644
--- a/net/test/cstruct.py
+++ b/net/test/cstruct.py
@@ -17,7 +17,8 @@
 Example usage:
 
 >>> # Declare a struct type by specifying name, field formats and field names.
-... # Field formats are the same as those used in the struct module.
+... # Field formats are the same as those used in the struct module, except:
+... # - S: Nested Struct.
 ... import cstruct
 >>> NLMsgHdr = cstruct.Struct("NLMsgHdr", "=LHHLL", "length type flags seq pid")
 >>>
@@ -44,6 +45,15 @@ NLMsgHdr(length=44, type=33, flags=2, seq=0, pid=510)
 >>> cstruct.Read(data, NLMsgHdr)
 (NLMsgHdr(length=44, type=33, flags=2, seq=0, pid=510), 'more data')
 >>>
+>>> # Structs can contain one or more nested structs. The nested struct types
+... # are specified in a list as an optional last argument. Nested structs may
+... # contain nested structs.
+... S = cstruct.Struct("S", "=BI", "byte1 int2")
+>>> N = cstruct.Struct("N", "!BSiS", "byte1 s2 int3 s2", [S, S])
+>>> NN = cstruct.Struct("NN", "SHS", "s1 word2 n3", [S, N])
+>>> nn = NN((S((1, 25000)), -29876, N((55, S((5, 6)), 1111, S((7, 8))))))
+>>> nn.n3.s2.int2 = 5
+>>>
 """
 
 import ctypes
@@ -52,9 +62,12 @@ import struct
 
 
 def CalcNumElements(fmt):
+  prevlen = len(fmt)
+  fmt = fmt.replace("S", "")
+  numstructs = prevlen - len(fmt)
   size = struct.calcsize(fmt)
   elements = struct.unpack(fmt, "\x00" * size)
-  return len(elements)
+  return len(elements) + numstructs
 
 
 def Struct(name, fmt, fieldnames, substructs={}):
diff --git a/net/test/cstruct_test.py b/net/test/cstruct_test.py
index 2d5a408..70ba90c 100755
--- a/net/test/cstruct_test.py
+++ b/net/test/cstruct_test.py
@@ -55,6 +55,45 @@ class CstructTest(unittest.TestCase):
       self.CheckEquals(i, i)
     self.CheckEquals(a1, a3)
 
+  def testNestedStructs(self):
+    Nested = cstruct.Struct("Nested", "!HSSi",
+                            "word1 nest2 nest3 int4",
+                            [TestStructA, TestStructB])
+    DoubleNested = cstruct.Struct("DoubleNested", "SSB",
+                                  "nest1 nest2 byte3",
+                                  [TestStructA, Nested])
+    d = DoubleNested((TestStructA((1, 2)),
+                      Nested((5, TestStructA((3, 4)), TestStructB((7, 8)), 9)),
+                      6))
+
+    expectedlen = (len(TestStructA) +
+                   2 + len(TestStructA) + len(TestStructB) + 4 +
+                   1)
+    self.assertEquals(expectedlen, len(DoubleNested))
+
+    self.assertEquals(7, d.nest2.nest3.byte1)
+
+    d.byte3 = 252
+    d.nest2.word1 = 33214
+    n = d.nest2
+    n.int4 = -55
+    t = n.nest3
+    t.int2 = 33627591
+
+    self.assertEquals(33627591, d.nest2.nest3.int2)
+
+    expected = (
+        "DoubleNested(nest1=TestStructA(byte1=1, int2=2),"
+        " nest2=Nested(word1=33214, nest2=TestStructA(byte1=3, int2=4),"
+        " nest3=TestStructB(byte1=7, int2=33627591), int4=-55), byte3=252)")
+    self.assertEquals(expected, str(d))
+    expected = ("01" "02000000"
+                "81be" "03" "04000000"
+                "07" "c71d0102" "ffffffc9" "fc").decode("hex")
+    self.assertEquals(expected, d.Pack())
+    unpacked = DoubleNested(expected)
+    self.CheckEquals(unpacked, d)
+
 
 if __name__ == "__main__":
   unittest.main()