From a4741e2bf2bb82cb6154935e481ba6db3c0094e2 Mon Sep 17 00:00:00 2001
From: Tim Rid <6593626+timrid@users.noreply.github.com>
Date: Sat, 20 Feb 2021 21:42:39 +0100
Subject: [PATCH] new setting in TStruct: "add_offsets"

---
 construct_typed/tstruct.py |  9 ++++++---
 tests/test_typed.py        | 24 ++++++++++++++++++++++--
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/construct_typed/tstruct.py b/construct_typed/tstruct.py
index 16f06ce..088ebd3 100644
--- a/construct_typed/tstruct.py
+++ b/construct_typed/tstruct.py
@@ -87,7 +87,7 @@ class _TStruct(Adapter[t.Any, t.Any, ContainerType, BuildTypes]):
     """
 
     def __init__(
-        self, container_type: t.Type[ContainerType], swapped: bool = False
+        self, container_type: t.Type[ContainerType], swapped: bool = False, add_offsets: bool = False
     ) -> None:
         if not issubclass(container_type, TContainerBase):
             raise TypeError(
@@ -101,6 +101,7 @@ class _TStruct(Adapter[t.Any, t.Any, ContainerType, BuildTypes]):
             )
         self.container_type = container_type
         self.swapped = swapped
+        self.add_offsets = add_offsets
 
         # get all fields from the dataclass
         fields = dataclasses.fields(self.container_type)
@@ -110,6 +111,8 @@ class _TStruct(Adapter[t.Any, t.Any, ContainerType, BuildTypes]):
         # extract the construct formats from the struct_type
         subcon_fields = {}
         for field in fields:
+            if add_offsets:
+                subcon_fields[f"@{field.name}"] = cs.Tell
             subcon_fields[field.name] = field.metadata["subcon"]
 
         # init adatper
@@ -175,7 +178,7 @@ class TStruct(_TStruct[ContainerType, ContainerType]):
     if t.TYPE_CHECKING:
 
         def __new__(
-            cls, container_type: t.Type[ContainerType], swapped: bool = False
+            cls, container_type: t.Type[ContainerType], swapped: bool = False, add_offsets: bool = False
         ) -> "TStruct[ContainerType]":
             ...
 
@@ -193,7 +196,7 @@ class TBitStruct(_TStruct[ContainerType, ContainerType]):
     if t.TYPE_CHECKING:
 
         def __new__(
-            cls, container_type: t.Type[ContainerType], swapped: bool = False
+            cls, container_type: t.Type[ContainerType], swapped: bool = False, add_offsets: bool = False
         ) -> "TBitStruct[ContainerType]":
             ...
 
diff --git a/tests/test_typed.py b/tests/test_typed.py
index be85227..e834fbe 100644
--- a/tests/test_typed.py
+++ b/tests/test_typed.py
@@ -108,6 +108,24 @@ def test_tstruct_swapped() -> None:
     assert str(normal.parse(b"\x00\x01\x02")) == str(swapped.parse(b"\x02\x00\x01"))
 
 
+def test_tstruct_add_offsets() -> None:
+    @dataclasses.dataclass
+    class TestContainer(cst.TContainerBase):
+        a: int = cst.TStructField(cs.Int16ub)
+        b: int = cst.TStructField(cs.Int8ub)
+
+    common(
+        cst.TStruct(TestContainer, add_offsets=True),
+        b"\x00\x01\x02",
+        TestContainer(a=1, b=2),
+        3,
+    )
+    c = cst.TStruct(TestContainer, add_offsets=True)
+    obj = c.parse(b"\x00\x01\x02")
+    assert obj["@a"] == 0
+    assert obj["@b"] == 2
+
+
 def test_tstruct_nested() -> None:
     @dataclasses.dataclass
     class TestContainer(cst.TContainerBase):
@@ -247,8 +265,10 @@ def test_tstruct_doc() -> None:
 
     assert format.subcon.a.docs == "This is the documentation of a"
     assert format.subcon.b.docs == "This is the documentation of b\nwhich is multiline"
-    assert format.subcon.c.docs == "This is the documentation of c\nwhich is also multiline"
-    
+    assert (
+        format.subcon.c.docs
+        == "This is the documentation of c\nwhich is also multiline"
+    )
 
 
 @pytest.mark.xfail(reason="not implemented yet")