From 976e68d4b9a3f4167b4056decf158cea55d5e25d Mon Sep 17 00:00:00 2001
From: Tim Rid <6593626+timrid@users.noreply.github.com>
Date: Tue, 13 Apr 2021 21:33:11 +0200
Subject: [PATCH] values from "cs.Const" and "cs.Default" are now set as
 default value from the dataclass

---
 construct_typed/tstruct.py | 21 +++++++++++++++++----
 tests/test_typed.py        | 27 +++++++++++++++++++++------
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/construct_typed/tstruct.py b/construct_typed/tstruct.py
index 4e18fa6..d413b33 100644
--- a/construct_typed/tstruct.py
+++ b/construct_typed/tstruct.py
@@ -61,18 +61,31 @@ def sfield(
     """
     Create a dataclass field for a "TStruct" and "TBitStruct" from a subcon.
     """
+    orig_subcon = subcon
+
     # Rename subcon, if doc or parsed are available
     if (doc is not None) or (parsed is not None):
         if doc is not None:
             doc = textwrap.dedent(doc).strip("\n")
         subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed)
 
-    if subcon.flagbuildnone is True:
-        # if subcon builds from "None", set default to "None"
+    if isinstance(orig_subcon, (cs.Const, cs.Default)) and not callable(
+        orig_subcon.value
+    ):
+        # Set simple values if cs.Const or cs.Default.
+        # If the value is a callable (or context lambda), then we cant set it because it
+        # is only defined at parsing/building
+        field = dataclasses.field(
+            default=orig_subcon.value,
+            init=False,
+            metadata={"subcon": subcon},
+        )
+    elif orig_subcon.flagbuildnone is True:
+        # If subcon builds from "None", set default to "None"
         field = dataclasses.field(
             default=None,
             init=False,
-            metadata={"subcon": cs.Renamed(subcon, newdocs=doc)},
+            metadata={"subcon": subcon},
         )
     else:
         field = dataclasses.field(metadata={"subcon": subcon})
@@ -170,7 +183,7 @@ class _TStruct(Adapter[t.Any, t.Any, ContainerType, BuildTypes]):
             fields = dataclasses.fields(self.container_type)
 
             # extract all fields from the container, that are used for create the dataclass object
-            ret_dict = {}
+            ret_dict: t.Dict[str, t.Any] = {}
             for field in fields:
                 value = getattr(obj, field.name)
                 ret_dict[field.name] = value
diff --git a/tests/test_typed.py b/tests/test_typed.py
index ab16b2f..be58c5e 100644
--- a/tests/test_typed.py
+++ b/tests/test_typed.py
@@ -11,6 +11,19 @@ import typing as t
 from .declarativeunittest import common, raises, setattrs
 
 
+def test_tcontainer_const_default() -> None:
+    @dataclasses.dataclass
+    class ConstDefaultTest(cst.TContainerMixin):
+        const_bytes: bytes = cst.sfield(cs.Const(b"BMP"))
+        const_int: int = cst.sfield(cs.Const(5, cs.Int8ub))
+        default_int: int = cst.sfield(cs.Default(cs.Int8ub, 28))
+
+    a = ConstDefaultTest()
+    assert a.const_bytes == b"BMP"
+    assert a.const_int == 5
+    assert a.default_int == 28
+
+
 def test_tcontainer_compare_with_dataclass() -> None:
     @dataclasses.dataclass
     class TestContainer:
@@ -27,7 +40,7 @@ def test_tcontainer_compare_with_dataclass() -> None:
 
     # ##### compare dot & dict access #####
     # dataclass
-    assert datacls.a == None
+    assert datacls.a == 1
     assert raises(lambda: datacls["a"]) == TypeError  # type: ignore
     assert datacls.b == 1
     assert raises(lambda: datacls["b"]) == TypeError  # type: ignore
@@ -41,8 +54,8 @@ def test_tcontainer_compare_with_dataclass() -> None:
         assert e.__class__ == TypeError
 
     # tcontainer
-    assert tcontainer.a == None
-    assert tcontainer["a"] == None
+    assert tcontainer.a == 1
+    assert tcontainer["a"] == 1
     assert tcontainer.b == 1
     assert tcontainer["b"] == 1
 
@@ -74,7 +87,8 @@ def test_tcontainer_order() -> None:
     format = cst.TStruct(Image)
     obj = Image(width=3, height=2)
     assert (
-        str(obj) == "Container: \n    signature = None\n    width = 3\n    height = 2"
+        str(obj)
+        == "Container: \n    signature = b'BMP' (total 3)\n    width = 3\n    height = 2"
     )
     obj = format.parse(format.build(obj))
     assert (
@@ -90,7 +104,7 @@ def test_tstruct() -> None:
         b: int = cst.sfield(cs.Int8ub)
 
     common(cst.TStruct(TestContainer), b"\x00\x01\x02", TestContainer(a=1, b=2), 3)
-
+    
     # check __getattr__
     c = cst.TStruct(TestContainer)
     assert c.a.name == "a"
@@ -193,6 +207,7 @@ def test_tstruct_const_field() -> None:
         == cs.ConstError
     )
 
+
 def test_tstruct_array_field() -> None:
     @dataclasses.dataclass
     class TestContainer(cst.TContainerMixin):
@@ -201,7 +216,7 @@ def test_tstruct_array_field() -> None:
     common(
         cst.TStruct(TestContainer),
         bytes(5),
-        TestContainer(array_field=[0,0,0,0,0]),
+        TestContainer(array_field=[0, 0, 0, 0, 0]),
         5,
     )