Add attribute offset support to cstruct am: 9c67278fa7 am: 5318026672 am: 97ffc5dc5d

am: 70175499c1

Change-Id: I06520e697fb14e326c31b560d26e91fbd38e82c3

2 files changed, 43 insertions, 2 deletions

diff --git a/net/test/cstruct.py b/net/test/cstruct.py
index d31979e..434552e 100644
--- a/net/test/cstruct.py
+++ b/net/test/cstruct.py
@@ -112,8 +112,7 @@ def Struct(name, fmt, fieldnames, substructs={}):
     # List of string fields that are ASCII strings.
     _asciiz = set()
 
-    if isinstance(_fieldnames, str):
-      _fieldnames = _fieldnames.split(" ")
+    _fieldnames = _fieldnames.split(" ")
 
     # Parse fmt into _format, converting any S format characters to "XXs",
     # where XX is the length of the struct type's packed representation.
@@ -137,6 +136,17 @@ def Struct(name, fmt, fieldnames, substructs={}):
 
     _length = CalcSize(_format)
 
+    offset_list = [0]
+    last_offset = 0
+    for i in xrange(len(_format)):
+      offset = CalcSize(_format[:i])
+      if offset > last_offset:
+        last_offset = offset
+        offset_list.append(offset)
+
+    # A dictionary that maps field names to their offsets in the struct.
+    _offsets = dict(zip(_fieldnames, offset_list))
+
     def _SetValues(self, values):
       # Replace self._values with the given list. We can't do direct assignment
       # because of the __setattr__ overload on this class.
@@ -197,6 +207,11 @@ def Struct(name, fmt, fieldnames, substructs={}):
       # callers get an unhelpful exception when they call Pack().
       self._values[self._FieldIndex(name)] = value
 
+    def offset(self, name):
+      if "." in name:
+        raise NotImplementedError("offset() on nested field")
+      return self._offsets[name]
+
     @classmethod
     def __len__(cls):
       return cls._length
diff --git a/net/test/cstruct_test.py b/net/test/cstruct_test.py
index b69aeb7..3e46af9 100755
--- a/net/test/cstruct_test.py
+++ b/net/test/cstruct_test.py
@@ -137,6 +137,32 @@ class CstructTest(unittest.TestCase):
     expected = ("00" + text_bytes + "00000000" "3412").decode("hex")
     self.assertEquals(expected, t1.Pack())
 
+  def testCstructOffset(self):
+    TestStruct = cstruct.Struct("TestStruct", "B16si16AH",
+                                "byte1 string2 int3 ascii4 word5")
+    nullstr = "hello" + (16 - len("hello")) * "\x00"
+    t = TestStruct((2, nullstr, 12345, nullstr, 33210))
+    self.assertEquals(0, t.offset("byte1"))
+    self.assertEquals(1, t.offset("string2"))  # sizeof(byte)
+    self.assertEquals(17, t.offset("int3"))    # sizeof(byte) + 16*sizeof(char)
+    # The integer is automatically padded by the struct module
+    # to match native alignment.
+    # offset = sizeof(byte) + 16*sizeof(char) + padding + sizeof(int)
+    self.assertEquals(24, t.offset("ascii4"))
+    self.assertEquals(40, t.offset("word5"))
+    self.assertRaises(KeyError, t.offset, "random")
+
+    # TODO: Add support for nested struct offset
+    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))
+    self.assertEqual(0, d.offset("nest1"))
+    self.assertEqual(len(TestStructA), d.offset("nest2"))
+    self.assertEqual(len(TestStructA) + len(Nested), d.offset("byte3"))
+    self.assertRaises(KeyError, t.offset, "word1")
 
 if __name__ == "__main__":
   unittest.main()