simplified "IfThenElse"

timrid · Jul 24, 2023 · 86fddbe · 86fddbe
1 parent 26f7fd8
commit 86fddbe
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 11 deletions.
diff --git 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
@@ -72,6 +72,15 @@ class Image(DataclassMixin):
         == "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