From 86fddbe2ace94e3d12e3455709aa06cbbc4ee6ed Mon Sep 17 00:00:00 2001
From: Tim Riddermann <tim.riddermann@winkhaus.de>
Date: Mon, 24 Jul 2023 10:14:52 +0200
Subject: [PATCH] simplified "IfThenElse"

---
 construct-stubs/core.pyi | 26 +++++++++++++++-----------
 tests/test_typed.py      |  9 +++++++++
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/construct-stubs/core.pyi b/construct-stubs/core.pyi
index aa1c4c7..23d9c43 100644
--- a/construct-stubs/core.pyi
+++ b/construct-stubs/core.pyi
@@ -749,25 +749,29 @@ ThenBuildTypes = t.TypeVar("ThenBuildTypes")
 ElseParsedType = t.TypeVar("ElseParsedType")
 ElseBuildTypes = t.TypeVar("ElseBuildTypes")
 
-class IfThenElse(
-    Construct[
-        t.Union[ThenParsedType, ElseParsedType], t.Union[ThenBuildTypes, ElseBuildTypes]
-    ]
-):
+class IfThenElse(Construct[ParsedType, BuildTypes]):
     condfunc: ConstantOrContextLambda[bool]
-    thensubcon: Construct[ThenParsedType, ThenBuildTypes]
-    elsesubcon: Construct[ElseParsedType, ElseBuildTypes]
-    def __init__(
-        self,
+    thensubcon: Construct[t.Any, t.Any]
+    elsesubcon: Construct[t.Any, t.Any]
+    @t.overload
+    def __new__(
+        cls: "type[IfThenElse[t.Union[ThenParsedType, ElseParsedType], t.Union[ThenBuildTypes, ElseBuildTypes]]]",
         condfunc: ConstantOrContextLambda[bool],
         thensubcon: Construct[ThenParsedType, ThenBuildTypes],
         elsesubcon: Construct[ElseParsedType, ElseBuildTypes],
-    ) -> None: ...
+    ) -> "IfThenElse[t.Union[ThenParsedType, ElseParsedType], t.Union[ThenBuildTypes, ElseBuildTypes]]": ...
+    @t.overload
+    def __new__(
+        cls: "type[IfThenElse[t.Any, t.Any]]",
+        condfunc: ConstantOrContextLambda[bool],
+        thensubcon: Construct[t.Any, t.Any],
+        elsesubcon: Construct[t.Any, t.Any],
+    ) -> "IfThenElse[t.Any, t.Any]": ...
 
 def If(
     condfunc: ConstantOrContextLambda[bool],
     subcon: Construct[ThenParsedType, ThenBuildTypes],
-) -> IfThenElse[ThenParsedType, None, ThenBuildTypes, None]: ...
+) -> IfThenElse[t.Optional[ThenParsedType], t.Optional[ThenBuildTypes]]: ...
 
 SwitchType = t.TypeVar("SwitchType")
 
diff --git a/tests/test_typed.py b/tests/test_typed.py
index 7d726a3..b85058a 100644
--- a/tests/test_typed.py
+++ b/tests/test_typed.py
@@ -72,6 +72,15 @@ def test_dataclass_str_repr() -> None:
         == "Image: \n    signature = b'BMP' (total 3)\n    width = 3\n    height = 2"
     )
 
+def test_dataclass_ifthenelse() -> None:
+    @dataclasses.dataclass
+    class IfThenElseTest(DataclassMixin):
+        test_if: t.Optional[int] = csfield(cs.If(False, cs.Int8ub))
+        test_ifthenelse: t.Optional[int] = csfield(cs.IfThenElse(True, cs.Int8ub, cs.Pass))
+
+    a = IfThenElseTest(test_if=None, test_ifthenelse=None)
+    assert a.test_if == None
+    assert a.test_ifthenelse == None
 
 def test_dataclass_struct() -> None:
     @dataclasses.dataclass