From ede0840a0b39f58e33a8973200b7a651a1596842 Mon Sep 17 00:00:00 2001 From: Tim Rid <6593626+timrid@users.noreply.github.com> Date: Fri, 1 Jan 2021 13:55:19 +0100 Subject: [PATCH] Fixed TStruct in TArray --- construct-stubs/core.pyi | 8 ++++---- construct_typed/__init__.py | 8 ++++---- construct_typed/generic_wrapper.py | 2 -- construct_typed/tarray.py | 25 +++++++------------------ construct_typed/tenum.py | 2 ++ construct_typed/tstruct.py | 28 +++++++--------------------- construct_typed/tunion.py | 2 ++ 7 files changed, 26 insertions(+), 49 deletions(-) diff --git a/construct-stubs/core.pyi b/construct-stubs/core.pyi index 916bfda..94cda69 100644 --- a/construct-stubs/core.pyi +++ b/construct-stubs/core.pyi @@ -168,9 +168,9 @@ class Subconstruct( class Adapter( Subconstruct[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes], ): - def __new__( - cls, subcon: Construct[SubconParsedType, SubconBuildTypes] - ) -> Adapter[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes]: ... + def __init__( + self, subcon: Construct[SubconParsedType, SubconBuildTypes] + ) -> None: ... def _decode( self, obj: SubconParsedType, context: Context, path: PathType ) -> ParsedType: ... @@ -503,7 +503,7 @@ class Default(Subconstruct[SubconParsedType, SubconBuildTypes, ParsedType, Build def __new__( cls, subcon: Construct[SubconParsedType, SubconBuildTypes], - value: SubconBuildTypes, + value: ConstantOrContextLambda2[SubconBuildTypes], ) -> Default[ SubconParsedType, SubconBuildTypes, diff --git a/construct_typed/__init__.py b/construct_typed/__init__.py index 5b894ab..a18f63d 100644 --- a/construct_typed/__init__.py +++ b/construct_typed/__init__.py @@ -1,7 +1,7 @@ -from .tarray import * -from .tenum import * -from .tstruct import * -from .tunion import * +from .tarray import TArray +from .tenum import TEnum +from .tstruct import TStruct, TBitStruct, StructField +from .tunion import TUnion, UnionField __all__ = [ "StructField", diff --git a/construct_typed/generic_wrapper.py b/construct_typed/generic_wrapper.py index a773cd4..cae9a51 100644 --- a/construct_typed/generic_wrapper.py +++ b/construct_typed/generic_wrapper.py @@ -8,10 +8,8 @@ ParsedType = t.TypeVar("ParsedType") BuildTypes = t.TypeVar("BuildTypes") SubconParsedType = t.TypeVar("SubconParsedType") SubconBuildTypes = t.TypeVar("SubconBuildTypes") -DataclassType = t.TypeVar("DataclassType") ListType = t.TypeVar("ListType") ValueType = t.TypeVar("ValueType") -EnumType = t.TypeVar("EnumType", bound=enum.IntEnum) if t.TYPE_CHECKING: diff --git a/construct_typed/tarray.py b/construct_typed/tarray.py index ced182a..f0929c8 100644 --- a/construct_typed/tarray.py +++ b/construct_typed/tarray.py @@ -5,31 +5,20 @@ from .generic_wrapper import * class TArray( Adapter[ - t.Any, - t.Any, - ParsedType, - BuildTypes, + SubconParsedType, + SubconBuildTypes, + t.List[SubconParsedType], + t.List[SubconParsedType], ] ): """ Adapter for an Array, that transforms the "ListContainer" to an standard "list" while parsing """ - # this is unfortunately needed because the stubs are using __new__ instead of __init__ - if t.TYPE_CHECKING: - - def __new__( - cls, - count: ConstantOrContextLambda[int], - subcon: Construct[SubconParsedType, SubconBuildTypes], - discard: bool = False, - ) -> "TArray[t.List[SubconParsedType], t.List[SubconParsedType]]": - ... - def __init__( self, count: ConstantOrContextLambda[int], - subcon: Construct[ParsedType, BuildTypes], + subcon: Construct[SubconParsedType, SubconBuildTypes], discard: bool = False, ) -> None: # init adatper @@ -42,8 +31,8 @@ class TArray( def _encode( self, - obj: BuildTypes, + obj: t.Any, context: Context, path: PathType, - ) -> t.List[BuildTypes]: + ) -> t.List[t.Any]: return obj # type: ignore diff --git a/construct_typed/tenum.py b/construct_typed/tenum.py index 8e2b718..315bcfe 100644 --- a/construct_typed/tenum.py +++ b/construct_typed/tenum.py @@ -3,6 +3,8 @@ import typing as t import construct as cs from .generic_wrapper import * +EnumType = t.TypeVar("EnumType", bound=enum.IntEnum) + class TEnum(Adapter[int, int, EnumType, t.Union[int, str, EnumType]]): def __new__( diff --git a/construct_typed/tstruct.py b/construct_typed/tstruct.py index bc4c751..eea8e88 100644 --- a/construct_typed/tstruct.py +++ b/construct_typed/tstruct.py @@ -3,6 +3,8 @@ import textwrap import dataclasses from .generic_wrapper import * +DataclassType = t.TypeVar("DataclassType") + def StructField( subcon: Construct[ParsedType, BuildTypes], @@ -44,7 +46,7 @@ class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): """ def __init__( - self, dataclass_type: t.Type[DataclassType], swapped: bool = False + self, dataclass_type: t.Type[ParsedType], swapped: bool = False ) -> None: if not dataclasses.is_dataclass(dataclass_type): raise TypeError( @@ -75,7 +77,7 @@ class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): def _decode( self, obj: "cs.Container[t.Any]", context: "cs.Context", path: "cs.PathType" - ) -> DataclassType: + ) -> ParsedType: # get all fields from the dataclass fields = dataclasses.fields(self.dataclass_type) @@ -98,7 +100,7 @@ class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): return dc def _encode( - self, obj: DataclassType, context: "cs.Context", path: "cs.PathType" + self, obj: ParsedType, context: "cs.Context", path: "cs.PathType" ) -> t.Dict[str, t.Any]: # get all fields from the dataclass fields = dataclasses.fields(self.dataclass_type) @@ -112,38 +114,22 @@ class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]): return ret_dict -class TStruct(_TStruct[ParsedType, BuildTypes]): +class TStruct(_TStruct[ParsedType, ParsedType]): """ Typed struct, based on standard dataclasses. """ - # this is unfortunately needed because the stubs are using __new__ instead of __init__ - if t.TYPE_CHECKING: - - def __new__( - cls, dataclass_type: t.Type[DataclassType], swapped: bool = False - ) -> "TStruct[DataclassType, DataclassType]": - ... - def _create_subcon( self, subcon_fields: t.Dict[str, t.Any] ) -> Construct[t.Any, t.Any]: return cs.Struct(**subcon_fields) -class TBitStruct(_TStruct[ParsedType, BuildTypes]): +class TBitStruct(_TStruct[ParsedType, ParsedType]): """ Typed bit struct, based on standard dataclasses. """ - # this is unfortunately needed because the stubs are using __new__ instead of __init__ - if t.TYPE_CHECKING: - - def __new__( - cls, dataclass_type: t.Type[DataclassType], swapped: bool = False - ) -> "TBitStruct[DataclassType, DataclassType]": - ... - def _create_subcon( self, subcon_fields: t.Dict[str, t.Any] ) -> Construct[t.Any, t.Any]: diff --git a/construct_typed/tunion.py b/construct_typed/tunion.py index 4c685a4..234faed 100644 --- a/construct_typed/tunion.py +++ b/construct_typed/tunion.py @@ -5,6 +5,8 @@ import construct as cs import dataclasses from .generic_wrapper import * +DataclassType = t.TypeVar("DataclassType") + def UnionField( subcon: Construct[ParsedType, BuildTypes],