From c5d08f9b87d63c0a5f09d092a68481f283f884e8 Mon Sep 17 00:00:00 2001
From: Tim Rid <6593626+timrid@users.noreply.github.com>
Date: Sun, 20 Feb 2022 20:33:28 +0100
Subject: [PATCH] add docstring to construct `docs`

---
 construct_typed/dataclass_struct.py |  8 ++++
 construct_typed/tenum.py            | 13 ++++++-
 tests/test_typed.py                 | 57 ++++++++++++++++++++++-------
 3 files changed, 62 insertions(+), 16 deletions(-)

diff --git a/construct_typed/dataclass_struct.py b/construct_typed/dataclass_struct.py
index cee8c89..0b799ee 100644
--- a/construct_typed/dataclass_struct.py
+++ b/construct_typed/dataclass_struct.py
@@ -271,6 +271,11 @@ class DataclassStruct:
         if not isinstance(reverse_fields, bool):  # type: ignore
             raise ValueError("`reverse_fields` parameter has to be an `bool` object")
 
+        # get documentation before creating the dataclass
+        docs = ""
+        if cls.__doc__ is not None:
+            docs = textwrap.dedent(cls.__doc__).strip("\n")
+
         # create dataclass
         dataclasses.dataclass(cls, kw_only=True)  # type: ignore
 
@@ -279,6 +284,9 @@ class DataclassStruct:
         if not isinstance(dc_constr, cs.Construct):  # type: ignore
             raise ValueError("`constr` sould return a `Construct` object")
 
+        # save docs
+        dc_constr.docs = docs
+
         # save construct format and make the class compatible to `Constructable` protocol
         setattr(cls, "__constr__", lambda: dc_constr)
 
diff --git a/construct_typed/tenum.py b/construct_typed/tenum.py
index 83f1d8c..5cb9116 100644
--- a/construct_typed/tenum.py
+++ b/construct_typed/tenum.py
@@ -1,9 +1,10 @@
 import enum
+import textwrap
 import typing as t
 
 import construct as cs
 
-from .generic import *
+from construct_typed.generic import *
 
 T = t.TypeVar("T")
 
@@ -26,6 +27,11 @@ class _EnumMeta(enum.EnumMeta):
         __namespace: t.Dict[str, t.Any],
         **kwargs: t.Any,
     ) -> T:
+        # get documentation before creating the enum
+        docs = ""
+        if "__doc__" in __namespace:
+            docs = textwrap.dedent(__namespace["__doc__"]).strip("\n")
+
         # create new enum object
         cls: T = super().__new__(metacls, __name, __bases, __namespace)  # type: ignore
 
@@ -49,7 +55,10 @@ class _EnumMeta(enum.EnumMeta):
         elif TFlags in __bases:
             enum_constr = TFlagsConstruct(subcon, cls)  # type: ignore
         else:
-            enum_constr = None
+            raise TypeError("neither `TEnum` nor `TFlags` in bases")
+
+        # save documentation
+        enum_constr.docs = docs
 
         # save construct format and make the class compatible to `Constructable` protocol
         setattr(cls, "__constr__", lambda: enum_constr)  # type: ignore
diff --git a/tests/test_typed.py b/tests/test_typed.py
index 24e29a0..f3d51d6 100644
--- a/tests/test_typed.py
+++ b/tests/test_typed.py
@@ -277,28 +277,39 @@ def test_dataclass_struct_wrong_container() -> None:
 
 
 def test_dataclass_struct_doc() -> None:
-    class TestContainer(DataclassStruct):
-        a: int = csfield(cs.Int16ub, doc="This is the documentation of a")
-        b: int = csfield(
-            cs.Int8ub, doc="This is the documentation of b\nwhich is multiline"
-        )
+    class TestContainer1(DataclassStruct):
+        """
+        Documentation of TestContainer
+        """
+
+        a: int = csfield(cs.Int16ub, doc="This is the doc of a")
+        b: int = csfield(cs.Int8ub, doc="This is the doc of b\nwhich is multiline")
         c: int = csfield(
             cs.Int8ub,
             doc="""
-            This is the documentation of c
+            This is the doc of c
             which is also multiline
             """,
         )
 
-    format = TestContainer.__constr__()
-    common(format, b"\x00\x01\x02\x03", TestContainer(a=1, b=2, c=3), 4)
+    format1 = TestContainer1.__constr__()
+    common(format1, b"\x00\x01\x02\x03", TestContainer1(a=1, b=2, c=3), 4)
 
-    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 format1.docs == "Documentation of TestContainer"
+    assert format1.subcon.a.docs == "This is the doc of a"
+    assert format1.subcon.b.docs == "This is the doc of b\nwhich is multiline"
+    assert format1.subcon.c.docs == "This is the doc of c\nwhich is also multiline"
+
+    class TestContainer2(DataclassStruct):
+        a: int = csfield(cs.Int16ub)
+        b: int = csfield(cs.Int8ub)
+        c: int = csfield(cs.Int8ub)
+
+    format2 = TestContainer2.__constr__()
+    assert format2.docs == ""
+    assert format2.subcon.a.docs == ""
+    assert format2.subcon.b.docs == ""
+    assert format2.subcon.c.docs == ""
 
 
 def test_dataclass_bitwise() -> None:
@@ -367,6 +378,24 @@ def test_tenum() -> None:
     assert raises(d.build, 8) == TypeError
 
 
+def test_tenum_doc() -> None:
+    class TestEnum1(TEnum, subcon=cs.Byte):
+        """
+        TestEnum documentation
+        """
+
+        one = 1
+
+    d1 = constr(TestEnum1)
+    assert d1.docs == "TestEnum documentation"
+
+    class TestEnum2(TEnum, subcon=cs.Byte):
+        two = 2
+
+    d2 = constr(TestEnum2)
+    assert d2.docs == ""
+
+
 def test_tenum_in_dataclass_struct() -> None:
     class TestEnum(TEnum, subcon=cs.Int8ub):
         a = 1