From 8211bb388dcf7ebf31ca2706c859ce1ba5efaec5 Mon Sep 17 00:00:00 2001
From: Tim Rid <6593626+timrid@users.noreply.github.com>
Date: Tue, 15 Dec 2020 22:27:01 +0100
Subject: [PATCH] added stubs for Timestamp, HexDump, Union

---
 construct-stubs/core.pyi | 51 +++++++++++++++++++++++++++++++++++++++-
 construct-stubs/expr.pyi |  6 ++---
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/construct-stubs/core.pyi b/construct-stubs/core.pyi
index d3dd539..1b936ef 100644
--- a/construct-stubs/core.pyi
+++ b/construct-stubs/core.pyi
@@ -1,6 +1,7 @@
 import typing as t
 import enum
-from construct.lib import Container, ListContainer, HexDisplayedBytes, HexDisplayedDict, HexDisplayedInteger
+import arrow  # type: ignore
+from construct.lib import Container, ListContainer, HexDisplayedBytes, HexDisplayedDict, HexDisplayedInteger, HexDumpDisplayedBytes, HexDumpDisplayedDict
 # unfortunately, there are a few duplications with "typing", e.g. Union and Optional, which is why the t. prefix must be used everywhere
 
 # Some of the Constructs can be optimised when the following typing optimisations are available:
@@ -409,6 +410,27 @@ class NamedTuple(Adapter[SubconParsedType, SubconBuildTypes, t.Tuple[t.Any, ...]
     ) -> None: ...
 
 
+@t.overload
+def Timestamp(
+    subcon: Construct[int, int],
+    unit: t.Literal["msdos"],
+    epoch: t.Literal["msdos"]
+) -> Adapter[int, int, arrow.Arrow, arrow.Arrow]: ...
+
+@t.overload
+def Timestamp(
+    subcon: Construct[int, int],
+    unit: t.Union[int, float],
+    epoch: t.Union[int, arrow.Arrow]
+) -> Adapter[int, int, arrow.Arrow, arrow.Arrow]: ...
+
+@t.overload
+def Timestamp(
+    subcon: Construct[float, float],
+    unit: t.Union[int, float],
+    epoch: t.Union[int, arrow.Arrow]
+) -> Adapter[float, float, arrow.Arrow, arrow.Arrow]: ...
+
 
 K = t.TypeVar("K")
 V = t.TypeVar("V")
@@ -434,10 +456,37 @@ class Hex(Adapter[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes]):
         subcon: Construct[SubconParsedType, SubconBuildTypes]
     ) -> Hex[SubconParsedType, SubconBuildTypes, SubconParsedType, SubconBuildTypes]: ...
 
+class HexDump(Adapter[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes]):
+    @t.overload
+    def __new__(
+        cls,
+        subcon: Construct[bytes, BuildTypes]
+    ) -> HexDump[bytes, BuildTypes, HexDumpDisplayedBytes, BuildTypes]: ...
+    @t.overload
+    def __new__(
+        cls,
+        subcon: Construct[t.Dict[K, V], BuildTypes]
+    ) -> HexDump[t.Dict[K, V], BuildTypes, HexDumpDisplayedDict[K, V], BuildTypes]: ...
+    @t.overload
+    def __new__(
+        cls,
+        subcon: Construct[SubconParsedType, SubconBuildTypes]
+    ) -> HexDump[SubconParsedType, SubconBuildTypes, SubconParsedType, SubconBuildTypes]: ...
+
 
 #===============================================================================
 # conditional
 #===============================================================================
+# this can maybe made better when variadic generics are available 
+class Union(Construct[Container[t.Any], t.Dict[str, t.Any]]):
+    def __init__(
+        self,
+        parsefrom: t.Optional[ConstantOrContextLambda[t.Union[int, str]]],
+        *subcons: Construct[t.Any, t.Any],
+        **subconskw: Construct[t.Any, t.Any]
+    ) -> None: ...
+
+
 # this can maybe made better when variadic generics are available
 class Select(Construct[ParsedType, BuildTypes]):
     def __new__(
diff --git a/construct-stubs/expr.pyi b/construct-stubs/expr.pyi
index c63c0c1..5e4f704 100644
--- a/construct-stubs/expr.pyi
+++ b/construct-stubs/expr.pyi
@@ -74,17 +74,17 @@ class ExprMixin(object):
 
 class UniExpr(ExprMixin):
     def __init__(self, op: UniOperator, operand: t.Any) -> None: ...
-    def __call__(self, obj: Union[Context, dict[str, t.Any]], *args: t.Any) -> t.Any: ...
+    def __call__(self, obj: t.Union[Context, dict[str, t.Any]], *args: t.Any) -> t.Any: ...
 
 
 class BinExpr(ExprMixin):
     def __init__(self, op: BinOperator, lhs: t.Any, rhs: t.Any) -> None: ...
-    def __call__(self, obj: Union[Context, dict[str, t.Any]], *args: t.Any) -> t.Any: ...
+    def __call__(self, obj: t.Union[Context, dict[str, t.Any]], *args: t.Any) -> t.Any: ...
 
 
 class Path(ExprMixin):
     def __init__(self, name: str, field: t.Optional[str] = ..., parent: t.Optional[Path] = ...) -> None: ...
-    def __call__(self, obj: Union[Context, dict[str, t.Any]], *args: t.Any) -> t.Any: ...
+    def __call__(self, obj: t.Union[Context, dict[str, t.Any]], *args: t.Any) -> t.Any: ...
     def __getattr__(self, name: str) -> Path: ...
     def __getitem__(self, name: str) -> Path: ...