From f5a3565f3e13581a729d427504c98eef0099b6d4 Mon Sep 17 00:00:00 2001 From: Tim Rid <6593626+timrid@users.noreply.github.com> Date: Sat, 2 Jan 2021 23:35:45 +0100 Subject: [PATCH] added stubs for LazyStruct, LazyListContainer, LazyArray, LazyBound fixed bug in test_typed.py --- construct-stubs/core.pyi | 41 ++++++++++++++++++++++++++++++++++++++++ tests/test_core.py | 14 +++++++------- tests/test_typed.py | 14 +++++++------- 3 files changed, 55 insertions(+), 14 deletions(-) diff --git a/construct-stubs/core.pyi b/construct-stubs/core.pyi index 4ffbdd1..19fef0e 100644 --- a/construct-stubs/core.pyi +++ b/construct-stubs/core.pyi @@ -5,12 +5,14 @@ import typing as t import arrow # type: ignore from construct.lib import ( Container, + ContainerType, HexDisplayedBytes, HexDisplayedDict, HexDisplayedInteger, HexDumpDisplayedBytes, HexDumpDisplayedDict, ListContainer, + ListType, ) # unfortunately, there are a few duplications with "typing", e.g. Union and Optional, which is why the t. prefix must be used everywhere @@ -928,6 +930,45 @@ class Lazy( ] ): ... +class LazyContainer(t.Generic[ContainerType], dict[str, ContainerType]): + def __getattr__(self, name: str) -> ContainerType: ... + def __getitem__(self, index: t.Union[str, int]) -> ContainerType: ... + def keys(self) -> t.Iterator[str]: ... + def values(self) -> t.List[ContainerType]: ... + def items(self) -> t.List[t.Tuple[str, ContainerType]]: ... + +class LazyStruct(Construct[ParsedType, BuildTypes]): + def __new__( + cls, *subcons: Construct[t.Any, t.Any], **subconskw: Construct[t.Any, t.Any] + ) -> LazyStruct[LazyContainer[t.Any], t.Optional[t.Dict[str, t.Any]]]: ... + def __getattr__(self, name: str) -> t.Any: ... + +class LazyListContainer(t.List[ListType]): ... + +class LazyArray( + Subconstruct[ + SubconParsedType, + SubconBuildTypes, + ParsedType, + BuildTypes, + ] +): + def __new__( + cls, + count: ConstantOrContextLambda[int], + subcon: Construct[SubconParsedType, SubconBuildTypes], + ) -> LazyArray[ + SubconParsedType, + SubconBuildTypes, + ListContainer[SubconParsedType], + t.List[SubconBuildTypes], + ]: ... + +class LazyBound(Construct[ParsedType, BuildTypes]): + def __new__( + cls, subconfunc: t.Callable[[], Construct[ParsedType, BuildTypes]] + ) -> LazyBound[ParsedType, BuildTypes]: ... + # =============================================================================== # adapters and validators # =============================================================================== diff --git a/tests/test_core.py b/tests/test_core.py index f819362..9fa90f4 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1360,22 +1360,22 @@ def test_lazyarray() -> None: assert raises(d.sizeof) == SizeofError def test_lazybound() -> None: - d = LazyBound(lambda: Byte) - common(d, b"\x01", 1) + d1 = LazyBound(lambda: Byte) + common(d1, b"\x01", 1) - d = Struct( + d2 = Struct( "value" / Byte, - "next" / If(this.value > 0, LazyBound(lambda: d)), + "next" / If(this.value > 0, LazyBound(lambda: d2)), ) - common(d, b"\x05\x09\x00", Container(value=5)(next=Container(value=9)(next=Container(value=0)(next=None)))) + common(d2, b"\x05\x09\x00", Container(value=5)(next=Container(value=9)(next=Container(value=0)(next=None)))) - d = Struct( + d3 = Struct( "value" / Byte, "next" / GreedyBytes, ) data = b"\x05\x09\x00" while data: - x = d.parse(data) + x = d3.parse(data) data = x.next print(x) diff --git a/tests/test_typed.py b/tests/test_typed.py index 50f80db..f3ceb85 100644 --- a/tests/test_typed.py +++ b/tests/test_typed.py @@ -13,12 +13,12 @@ from .declarativeunittest import common, raises, setattrs def test_tcontainer_compare_with_dataclass() -> None: @dataclasses.dataclass class TestContainer: - a: cst.Optional[int] = cst.TStructField(cs.Const(1, cs.Byte)) + a: cst.Opt[int] = cst.TStructField(cs.Const(1, cs.Byte)) b: int = cst.TStructField(cs.Int8ub) @dataclasses.dataclass class TestTContainer(cst.TContainerBase): - a: cst.Optional[int] = cst.TStructField(cs.Const(1, cs.Byte)) + a: cst.Opt[int] = cst.TStructField(cs.Const(1, cs.Byte)) b: int = cst.TStructField(cs.Int8ub) datacls = TestContainer(b=1) @@ -66,7 +66,7 @@ def test_tcontainer_compare_with_dataclass() -> None: def test_tcontainer_order() -> None: @dataclasses.dataclass class Image(cst.TContainerBase): - signature: cst.Optional[bytes] = cst.TStructField(cs.Const(b"BMP")) + signature: cst.Opt[bytes] = cst.TStructField(cs.Const(b"BMP")) width: int = cst.TStructField(cs.Int8ub) height: int = cst.TStructField(cs.Int8ub) @@ -130,7 +130,7 @@ def test_tstruct_default_field() -> None: class Image(cst.TContainerBase): width: int = cst.TStructField(cs.Int8ub) height: int = cst.TStructField(cs.Int8ub) - pixels: cst.Optional[bytes] = cst.TStructField( + pixels: cst.Opt[bytes] = cst.TStructField( cs.Default( cs.Bytes(cs.this.width * cs.this.height), lambda ctx: bytes(ctx.width * ctx.height), @@ -148,7 +148,7 @@ def test_tstruct_default_field() -> None: def test_tstruct_const_field() -> None: @dataclasses.dataclass class TestContainer(cst.TContainerBase): - const_field: cst.Optional[bytes] = cst.TStructField(cs.Const(b"\x00")) + const_field: cst.Opt[bytes] = cst.TStructField(cs.Const(b"\x00")) common( cst.TStruct(TestContainer), @@ -169,7 +169,7 @@ def test_tstruct_const_field() -> None: def test_tstruct_anonymus_fields_1() -> None: @dataclasses.dataclass class TestContainer(cst.TContainerBase): - _1: cst.Optional[bytes] = cst.TStructField(cs.Const(b"\x00")) + _1: cst.Opt[bytes] = cst.TStructField(cs.Const(b"\x00")) _2: None = cst.TStructField(cs.Padding(1)) _3: None = cst.TStructField(cs.Pass) _4: None = cst.TStructField(cs.Terminated) @@ -186,7 +186,7 @@ def test_tstruct_anonymus_fields_2() -> None: @dataclasses.dataclass class TestContainer(cst.TContainerBase): _1: int = cst.TStructField(cs.Computed(7)) - _2: cst.Optional[bytes] = cst.TStructField(cs.Const(b"JPEG")) + _2: cst.Opt[bytes] = cst.TStructField(cs.Const(b"JPEG")) _3: None = cst.TStructField(cs.Pass) _4: None = cst.TStructField(cs.Terminated)