- divided construct_typed in multiple files
- adapted the "test_typed.py" to dataclasses
This commit is contained in:
Tim Rid
parent
aff5c0589b
commit
e6f3f12d30
10 changed files with 736 additions and 295 deletions
288
construct_typed/__init__ old.py
Normal file
288
construct_typed/__init__ old.py
Normal file
|
|
@ -0,0 +1,288 @@
|
|||
import enum
|
||||
import typing as t
|
||||
import textwrap
|
||||
import construct as cs
|
||||
import dataclasses
|
||||
import .generic_construct as gcs
|
||||
|
||||
ParsedType = t.TypeVar("ParsedType")
|
||||
BuildTypes = t.TypeVar("BuildTypes")
|
||||
SubconParsedType = t.TypeVar("SubconParsedType")
|
||||
SubconBuildTypes = t.TypeVar("SubconBuildTypes")
|
||||
EnumType = t.TypeVar("EnumType", bound=enum.IntEnum)
|
||||
DataclassType = t.TypeVar("DataclassType")
|
||||
ListType = t.TypeVar("ListType")
|
||||
|
||||
|
||||
# if t.TYPE_CHECKING:
|
||||
# # while type checking, the original classes are generics, because they are defined in the stubs.
|
||||
# from construct import Construct, Adapter
|
||||
# from construct import ListContainer as List
|
||||
# else:
|
||||
# # at runtime, the original classes are no generics, so whe have to make new classes with generics support
|
||||
# class Construct(t.Generic[ParsedType, BuildTypes], cs.Construct):
|
||||
# pass
|
||||
|
||||
# class Adapter(
|
||||
# t.Generic[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes],
|
||||
# cs.Adapter,
|
||||
# ):
|
||||
# pass
|
||||
|
||||
# class List(t.Generic[ListType], cs.ListContainer):
|
||||
# pass
|
||||
|
||||
|
||||
|
||||
# ===============================================================================
|
||||
# mappings
|
||||
# ===============================================================================
|
||||
class TEnum(gcs.Adapter[int, int, EnumType, t.Union[int, str, EnumType]]):
|
||||
def __new__(
|
||||
cls, subcon: gcs.Construct[int, int], enum_type: t.Type[EnumType]
|
||||
) -> "TEnum[EnumType]":
|
||||
return super(TEnum, cls).__new__(cls, subcon, enum_type) # type: ignore
|
||||
|
||||
def __init__(self, subcon: gcs.Construct[int, int], enum_type: t.Type[EnumType]):
|
||||
if not issubclass(enum_type, enum.IntEnum):
|
||||
raise TypeError(
|
||||
'The class "{}" is not an "enum.IntEnum"'.format(
|
||||
enum_type.__name__
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _missing_(
|
||||
cls: t.Type[EnumType], value: t.Union[int, EnumType]
|
||||
) -> t.Optional[EnumType]:
|
||||
if isinstance(value, int):
|
||||
return cls._create_pseudo_member_(value)
|
||||
return None # will raise the ValueError in Enum.__new__
|
||||
|
||||
@classmethod
|
||||
def _create_pseudo_member_(cls: t.Type[EnumType], value: int) -> EnumType:
|
||||
pseudo_member = cls._value2member_map_.get(value, None)
|
||||
if pseudo_member is None:
|
||||
new_member = int.__new__(cls, value)
|
||||
# I expect a name attribute to hold a string, hence str(value)
|
||||
# However, new_member._name_ = value works, too
|
||||
new_member._name_ = str(value)
|
||||
new_member._value_ = value
|
||||
pseudo_member = cls._value2member_map_.setdefault(value, new_member)
|
||||
return pseudo_member
|
||||
|
||||
# Monkey-patch the enum type with __missing__ method. So if a enum value
|
||||
# not found in the enum a new pseudo member is created.
|
||||
# The idea is taken from: https://stackoverflow.com/a/57179436
|
||||
enum_type._missing_ = _missing_
|
||||
enum_type._create_pseudo_member_ = _create_pseudo_member_
|
||||
|
||||
# save enum type
|
||||
self.enum_type = enum_type
|
||||
|
||||
# init adatper
|
||||
super(TEnum, self).__init__(subcon) # type: ignore
|
||||
|
||||
def _decode(self, obj: int, context: "cs.Context", path: "cs.PathType") -> EnumType:
|
||||
return self.enum_type(obj)
|
||||
|
||||
def _encode(
|
||||
self,
|
||||
obj: t.Union[int, str, EnumType],
|
||||
context: "cs.Context",
|
||||
path: "cs.PathType",
|
||||
) -> int:
|
||||
try:
|
||||
if isinstance(obj, str):
|
||||
return int(self.enum_type[obj])
|
||||
else:
|
||||
return int(self.enum_type(obj))
|
||||
except:
|
||||
raise cs.MappingError(
|
||||
"building failed, no mapping for %r" % (obj,), path=path
|
||||
)
|
||||
|
||||
|
||||
# ===============================================================================
|
||||
# structures and sequences
|
||||
# ===============================================================================
|
||||
def StructField(
|
||||
subcon: gcs.Construct[ParsedType, BuildTypes],
|
||||
doc: t.Optional[str] = None,
|
||||
parsed: t.Optional[t.Callable[[t.Any, "cs.Context"], None]] = None,
|
||||
) -> ParsedType:
|
||||
"""
|
||||
Create a dataclass field for a "TStruct" and "TBitStruct" from a subcon.
|
||||
"""
|
||||
# Rename subcon, if doc or parsed are available
|
||||
if (doc is not None) or (parsed is not None):
|
||||
if doc is not None:
|
||||
doc = textwrap.dedent(doc)
|
||||
subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed)
|
||||
|
||||
if subcon.flagbuildnone is True:
|
||||
# some subcons have a predefined default value. all other have "None"
|
||||
default: t.Any = None
|
||||
if isinstance(subcon, (cs.Const, cs.Default)):
|
||||
if callable(subcon.value):
|
||||
raise ValueError("lamda as default is not supported")
|
||||
default = subcon.value
|
||||
|
||||
# if subcon builds from "None", set default to "None"
|
||||
field = dataclasses.field(
|
||||
default=default,
|
||||
init=False,
|
||||
metadata={"subcon": cs.Renamed(subcon, newdocs=doc)},
|
||||
)
|
||||
else:
|
||||
field = dataclasses.field(metadata={"subcon": subcon})
|
||||
|
||||
return field # type: ignore
|
||||
|
||||
|
||||
class _TStruct(
|
||||
gcs.Adapter["cs.Container[t.Any]", t.Dict[str, t.Any], DataclassType, DataclassType]
|
||||
):
|
||||
"""
|
||||
Base class for a typed struct, based on standard dataclasses.
|
||||
"""
|
||||
|
||||
def __new__(
|
||||
cls, dataclass_type: t.Type[DataclassType], swapped: bool = False
|
||||
) -> "_TStruct[DataclassType]":
|
||||
return super(_TStruct, cls).__new__(cls, dataclass_type, swapped) # type: ignore
|
||||
|
||||
def __init__(
|
||||
self, dataclass_type: t.Type[DataclassType], swapped: bool = False
|
||||
) -> None:
|
||||
if not dataclasses.is_dataclass(dataclass_type):
|
||||
raise TypeError(
|
||||
'The class "{}" is not a "dataclasses.dataclass"'.format(
|
||||
dataclass_type.__name__
|
||||
)
|
||||
)
|
||||
self.dataclass_type = dataclass_type
|
||||
self.swapped = swapped
|
||||
|
||||
# get all fields from the dataclass
|
||||
fields = dataclasses.fields(self.dataclass_type)
|
||||
if self.swapped:
|
||||
fields = tuple(reversed(fields))
|
||||
|
||||
# extract the construct formats from the struct_type
|
||||
subcon_fields = {}
|
||||
for field in fields:
|
||||
subcon_fields[field.name] = field.metadata["subcon"]
|
||||
|
||||
# init adatper
|
||||
super(_TStruct, self).__init__(self._create_subcon(subcon_fields)) # type: ignore
|
||||
|
||||
def _create_subcon(
|
||||
self, subcon_fields: t.Dict[str, t.Any]
|
||||
) -> gcs.Construct[t.Any, t.Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
def _decode(
|
||||
self, obj: "cs.Container[t.Any]", context: "cs.Context", path: "cs.PathType"
|
||||
) -> DataclassType:
|
||||
# get all fields from the dataclass
|
||||
fields = dataclasses.fields(self.dataclass_type)
|
||||
|
||||
# extract all fields from the container, that are used for create the dataclass object
|
||||
dc_init = {}
|
||||
for field in fields:
|
||||
if field.init:
|
||||
value = getattr(obj, field.name)
|
||||
dc_init[field.name] = value
|
||||
|
||||
# create object of dataclass
|
||||
dc = self.dataclass_type(**dc_init) # type: ignore
|
||||
|
||||
# extract all other values from the container, an pass it to the dataclass
|
||||
for field in fields:
|
||||
if not field.init:
|
||||
value = getattr(obj, field.name)
|
||||
setattr(dc, field.name, value)
|
||||
|
||||
return dc
|
||||
|
||||
def _encode(
|
||||
self, obj: DataclassType, context: "cs.Context", path: "cs.PathType"
|
||||
) -> t.Dict[str, t.Any]:
|
||||
# get all fields from the dataclass
|
||||
fields = dataclasses.fields(self.dataclass_type)
|
||||
|
||||
# extract all fields from the container, that are used for create the dataclass object
|
||||
ret_dict = {}
|
||||
for field in fields:
|
||||
value = getattr(obj, field.name)
|
||||
ret_dict[field.name] = value
|
||||
|
||||
return ret_dict
|
||||
|
||||
|
||||
class TStruct(_TStruct[DataclassType]):
|
||||
"""
|
||||
Typed struct, based on standard dataclasses.
|
||||
"""
|
||||
|
||||
def _create_subcon(
|
||||
self, subcon_fields: t.Dict[str, t.Any]
|
||||
) -> gcs.Construct[t.Any, t.Any]:
|
||||
return cs.Struct(**subcon_fields)
|
||||
|
||||
|
||||
class TBitStruct(_TStruct[DataclassType]):
|
||||
"""
|
||||
Typed bit struct, based on standard dataclasses.
|
||||
"""
|
||||
|
||||
def _create_subcon(
|
||||
self, subcon_fields: t.Dict[str, t.Any]
|
||||
) -> gcs.Construct[t.Any, t.Any]:
|
||||
return cs.BitStruct(**subcon_fields)
|
||||
|
||||
|
||||
|
||||
#===============================================================================
|
||||
# conditional
|
||||
#===============================================================================
|
||||
def UnionField(
|
||||
subcon: gcs.Construct[ParsedType, BuildTypes],
|
||||
doc: t.Optional[str] = None,
|
||||
parsed: t.Optional[t.Callable[[t.Any, "cs.Context"], None]] = None,
|
||||
) -> ParsedType:
|
||||
"""
|
||||
Create a dataclass field for a "TUnion" from a subcon.
|
||||
"""
|
||||
# Rename subcon, if doc or parsed are available
|
||||
if (doc is not None) or (parsed is not None):
|
||||
if doc is not None:
|
||||
doc = textwrap.dedent(doc)
|
||||
subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed)
|
||||
|
||||
if subcon.flagbuildnone is True:
|
||||
# some subcons have a predefined default value. all other have "None"
|
||||
default: t.Any = None
|
||||
if isinstance(subcon, (cs.Const, cs.Default)):
|
||||
if callable(subcon.value):
|
||||
raise ValueError("lamda as default is not supported")
|
||||
default = subcon.value
|
||||
|
||||
# if subcon builds from "None", set default to "None"
|
||||
field = dataclasses.field(
|
||||
default=default,
|
||||
init=False,
|
||||
metadata={"subcon": cs.Renamed(subcon, newdocs=doc)},
|
||||
)
|
||||
else:
|
||||
field = dataclasses.field(metadata={"subcon": subcon})
|
||||
|
||||
return field # type: ignore
|
||||
|
||||
|
||||
|
||||
# TODO: TypedUnion
|
||||
# TODO: TypedLazyStruct
|
||||
# TODO: TypedSequence: Based on typing.namedtuple
|
||||
# TODO: FocusedSeq: Based on typing.namedtuple
|
||||
|
|
@ -1,235 +1,14 @@
|
|||
from enum import IntEnum
|
||||
import typing as t
|
||||
import textwrap
|
||||
import construct as cs
|
||||
import dataclasses
|
||||
from .tarray import *
|
||||
from .tenum import *
|
||||
from .tstruct import *
|
||||
from .tunion import *
|
||||
|
||||
ParsedType = t.TypeVar("ParsedType")
|
||||
BuildTypes = t.TypeVar("BuildTypes")
|
||||
SubconParsedType = t.TypeVar("SubconParsedType")
|
||||
SubconBuildTypes = t.TypeVar("SubconBuildTypes")
|
||||
EnumType = t.TypeVar("EnumType", bound=IntEnum)
|
||||
DataclassType = t.TypeVar("DataclassType")
|
||||
|
||||
|
||||
if t.TYPE_CHECKING:
|
||||
# while type checking, the original classes are generics, because they are defined in the stubs.
|
||||
from construct import Construct, Adapter
|
||||
else:
|
||||
# at runtime, the original classes are no generics, so whe have to make new classes with generics support
|
||||
class Construct(t.Generic[ParsedType, BuildTypes], cs.Construct):
|
||||
pass
|
||||
|
||||
class Adapter(
|
||||
t.Generic[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes],
|
||||
cs.Adapter,
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
# ===============================================================================
|
||||
# mappings
|
||||
# ===============================================================================
|
||||
class TEnum(Adapter[int, int, EnumType, t.Union[int, str, EnumType]]):
|
||||
def __new__(
|
||||
cls, subcon: Construct[int, int], enum_type: t.Type[EnumType]
|
||||
) -> "TEnum[EnumType]":
|
||||
return super(TEnum, cls).__new__(cls, subcon, enum_type) # type: ignore
|
||||
|
||||
def __init__(self, subcon: Construct[int, int], enum_type: t.Type[EnumType]):
|
||||
@classmethod
|
||||
def _missing_(
|
||||
cls: t.Type[EnumType], value: t.Union[int, EnumType]
|
||||
) -> t.Optional[EnumType]:
|
||||
if isinstance(value, int):
|
||||
return cls._create_pseudo_member_(value)
|
||||
return None # will raise the ValueError in Enum.__new__
|
||||
|
||||
@classmethod
|
||||
def _create_pseudo_member_(cls: t.Type[EnumType], value: int) -> EnumType:
|
||||
pseudo_member = cls._value2member_map_.get(value, None)
|
||||
if pseudo_member is None:
|
||||
new_member = int.__new__(cls, value)
|
||||
# I expect a name attribute to hold a string, hence str(value)
|
||||
# However, new_member._name_ = value works, too
|
||||
new_member._name_ = str(value)
|
||||
new_member._value_ = value
|
||||
pseudo_member = cls._value2member_map_.setdefault(value, new_member)
|
||||
return pseudo_member
|
||||
|
||||
# Monkey-patch the enum type with __missing__ method. So if a enum value
|
||||
# not found in the enum a new pseudo member is created.
|
||||
# The idea is taken from: https://stackoverflow.com/a/57179436
|
||||
enum_type._missing_ = _missing_
|
||||
enum_type._create_pseudo_member_ = _create_pseudo_member_
|
||||
|
||||
# save enum type
|
||||
self.enum_type = enum_type
|
||||
|
||||
# init adatper
|
||||
super(TEnum, self).__init__(subcon) # type: ignore
|
||||
|
||||
def _decode(self, obj: int, context: "cs.Context", path: "cs.PathType") -> EnumType:
|
||||
return self.enum_type(obj)
|
||||
|
||||
def _encode(
|
||||
self,
|
||||
obj: t.Union[int, str, EnumType],
|
||||
context: "cs.Context",
|
||||
path: "cs.PathType",
|
||||
) -> int:
|
||||
try:
|
||||
if isinstance(obj, str):
|
||||
return int(self.enum_type[obj])
|
||||
else:
|
||||
return int(self.enum_type(obj))
|
||||
except:
|
||||
raise cs.MappingError(
|
||||
"building failed, no mapping for %r" % (obj,), path=path
|
||||
)
|
||||
|
||||
|
||||
# ===============================================================================
|
||||
# structures and sequences
|
||||
# ===============================================================================
|
||||
def TSubcon(
|
||||
subcon: Construct[ParsedType, BuildTypes],
|
||||
doc: t.Optional[str] = None,
|
||||
parsed: t.Optional[t.Callable[[t.Any, "cs.Context"], None]] = None,
|
||||
) -> ParsedType:
|
||||
"""
|
||||
Create a dataclass field from a subcon.
|
||||
"""
|
||||
# Rename subcon, if doc or parsed are available
|
||||
if (doc is not None) or (parsed is not None):
|
||||
if doc is not None:
|
||||
doc = textwrap.dedent(doc)
|
||||
subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed)
|
||||
|
||||
if subcon.flagbuildnone is True:
|
||||
# some subcons have a predefined default value. all other have "None"
|
||||
default: t.Any = None
|
||||
if isinstance(subcon, (cs.Const, cs.Default)):
|
||||
if callable(subcon.value):
|
||||
raise ValueError("lamda as default is not supported")
|
||||
default = subcon.value
|
||||
|
||||
# if subcon builds from "None", set default to "None"
|
||||
field = dataclasses.field(
|
||||
default=default,
|
||||
init=False,
|
||||
metadata={"subcon": cs.Renamed(subcon, newdocs=doc)},
|
||||
)
|
||||
else:
|
||||
field = dataclasses.field(metadata={"subcon": subcon})
|
||||
|
||||
return field # type: ignore
|
||||
|
||||
|
||||
class _TStruct(
|
||||
Adapter["cs.Container[t.Any]", t.Dict[str, t.Any], DataclassType, DataclassType]
|
||||
):
|
||||
"""
|
||||
Base class for a typed struct, based on standard dataclasses.
|
||||
"""
|
||||
|
||||
def __new__(
|
||||
cls, dataclass_type: t.Type[DataclassType], swapped: bool = False
|
||||
) -> "_TStruct[DataclassType]":
|
||||
return super(_TStruct, cls).__new__(cls, dataclass_type, swapped) # type: ignore
|
||||
|
||||
def __init__(
|
||||
self, dataclass_type: t.Type[DataclassType], swapped: bool = False
|
||||
) -> None:
|
||||
if not dataclasses.is_dataclass(dataclass_type):
|
||||
raise TypeError(
|
||||
'the "dataclass_type" has to be a dataclass but is "{}"'.format(
|
||||
type(dataclass_type).__name__
|
||||
)
|
||||
)
|
||||
self.dataclass_type = dataclass_type
|
||||
self.swapped = swapped
|
||||
|
||||
# get all fields from the dataclass
|
||||
fields = dataclasses.fields(self.dataclass_type)
|
||||
if self.swapped:
|
||||
fields = tuple(reversed(fields))
|
||||
|
||||
# extract the construct formats from the struct_type
|
||||
subcon_fields = {}
|
||||
for field in fields:
|
||||
subcon_fields[field.name] = field.metadata["subcon"]
|
||||
|
||||
# init adatper
|
||||
super(_TStruct, self).__init__(self._create_subcon(subcon_fields)) # type: ignore
|
||||
|
||||
def _create_subcon(
|
||||
self, subcon_fields: t.Dict[str, t.Any]
|
||||
) -> Construct[t.Any, t.Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
def _decode(
|
||||
self, obj: "cs.Container[t.Any]", context: "cs.Context", path: "cs.PathType"
|
||||
) -> DataclassType:
|
||||
# get all fields from the dataclass
|
||||
fields = dataclasses.fields(self.dataclass_type)
|
||||
|
||||
# extract all fields from the container, that are used for create the dataclass object
|
||||
dc_init = {}
|
||||
for field in fields:
|
||||
if field.init:
|
||||
value = getattr(obj, field.name)
|
||||
dc_init[field.name] = value
|
||||
|
||||
# create object of dataclass
|
||||
dc = self.dataclass_type(**dc_init) # type: ignore
|
||||
|
||||
# extract all other values from the container, an pass it to the dataclass
|
||||
for field in fields:
|
||||
if not field.init:
|
||||
value = getattr(obj, field.name)
|
||||
setattr(dc, field.name, value)
|
||||
|
||||
return dc
|
||||
|
||||
def _encode(
|
||||
self, obj: DataclassType, context: "cs.Context", path: "cs.PathType"
|
||||
) -> t.Dict[str, t.Any]:
|
||||
# get all fields from the dataclass
|
||||
fields = dataclasses.fields(self.dataclass_type)
|
||||
|
||||
# extract all fields from the container, that are used for create the dataclass object
|
||||
ret_dict = {}
|
||||
for field in fields:
|
||||
value = getattr(obj, field.name)
|
||||
ret_dict[field.name] = value
|
||||
|
||||
return ret_dict
|
||||
|
||||
|
||||
class TStruct(_TStruct[DataclassType]):
|
||||
"""
|
||||
Typed struct, based on standard dataclasses.
|
||||
"""
|
||||
|
||||
def _create_subcon(
|
||||
self, subcon_fields: t.Dict[str, t.Any]
|
||||
) -> Construct[t.Any, t.Any]:
|
||||
return cs.Struct(**subcon_fields)
|
||||
|
||||
|
||||
class TBitStruct(_TStruct[DataclassType]):
|
||||
"""
|
||||
Typed bit struct, based on standard dataclasses.
|
||||
"""
|
||||
|
||||
def _create_subcon(
|
||||
self, subcon_fields: t.Dict[str, t.Any]
|
||||
) -> Construct[t.Any, t.Any]:
|
||||
return cs.BitStruct(**subcon_fields)
|
||||
|
||||
|
||||
# TODO: TypedUnion
|
||||
# TODO: TypedLazyStruct
|
||||
# TODO: TypedSequence: Based on typing.namedtuple
|
||||
# TODO: FocusedSeq: Based on typing.namedtuple
|
||||
__all__ = [
|
||||
"StructField",
|
||||
"TStruct",
|
||||
"TBitStruct",
|
||||
"TEnum",
|
||||
"UnionField",
|
||||
"TUnion",
|
||||
"TArray"
|
||||
]
|
||||
|
|
|
|||
43
construct_typed/generic_wrapper.py
Normal file
43
construct_typed/generic_wrapper.py
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
import enum
|
||||
import typing as t
|
||||
import textwrap
|
||||
import construct as cs
|
||||
import dataclasses
|
||||
|
||||
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:
|
||||
# while type checking, the original classes are generics, because they are defined in the stubs.
|
||||
from construct import Construct as Construct
|
||||
from construct import Adapter as Adapter
|
||||
from construct import ListContainer as ListContainer
|
||||
from construct import Context as Context
|
||||
from construct import ConstantOrContextLambda as ConstantOrContextLambda
|
||||
from construct import PathType as PathType
|
||||
else:
|
||||
# at runtime, the original classes are no generics, so whe have to make new classes with generics support
|
||||
class Construct(t.Generic[ParsedType, BuildTypes], cs.Construct):
|
||||
pass
|
||||
|
||||
class Adapter(
|
||||
t.Generic[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes],
|
||||
cs.Adapter,
|
||||
):
|
||||
pass
|
||||
|
||||
class ListContainer(t.Generic[ListType], cs.ListContainer):
|
||||
pass
|
||||
|
||||
class Context:
|
||||
pass
|
||||
|
||||
ConstantOrContextLambda = t.Union[ValueType, t.Callable[[Context], t.Any]]
|
||||
PathType = str
|
||||
0
construct_typed/py.typed
Normal file
0
construct_typed/py.typed
Normal file
49
construct_typed/tarray.py
Normal file
49
construct_typed/tarray.py
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
import typing as t
|
||||
import construct as cs
|
||||
from .generic_wrapper import *
|
||||
|
||||
|
||||
class TArray(
|
||||
Adapter[
|
||||
t.Any,
|
||||
t.Any,
|
||||
ParsedType,
|
||||
BuildTypes,
|
||||
]
|
||||
):
|
||||
"""
|
||||
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],
|
||||
discard: bool = False,
|
||||
) -> None:
|
||||
# init adatper
|
||||
super(TArray, self).__init__(cs.Array(count, subcon, discard)) # type: ignore
|
||||
|
||||
def _decode(
|
||||
self, obj: ListContainer[ParsedType], context: Context, path: PathType
|
||||
) -> ParsedType:
|
||||
return list(obj) # type: ignore
|
||||
|
||||
def _encode(
|
||||
self,
|
||||
obj: BuildTypes,
|
||||
context: Context,
|
||||
path: PathType,
|
||||
) -> t.List[BuildTypes]:
|
||||
return obj # type: ignore
|
||||
68
construct_typed/tenum.py
Normal file
68
construct_typed/tenum.py
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
import enum
|
||||
import typing as t
|
||||
import construct as cs
|
||||
from .generic_wrapper import *
|
||||
|
||||
|
||||
class TEnum(Adapter[int, int, EnumType, t.Union[int, str, EnumType]]):
|
||||
def __new__(
|
||||
cls, subcon: Construct[int, int], enum_type: t.Type[EnumType]
|
||||
) -> "TEnum[EnumType]":
|
||||
return super(TEnum, cls).__new__(cls, subcon, enum_type) # type: ignore
|
||||
|
||||
def __init__(self, subcon: Construct[int, int], enum_type: t.Type[EnumType]):
|
||||
if not issubclass(enum_type, enum.IntEnum):
|
||||
raise TypeError(
|
||||
'The class "{}" is not an "enum.IntEnum"'.format(enum_type.__name__)
|
||||
)
|
||||
|
||||
@classmethod # type: ignore
|
||||
def _missing_(
|
||||
cls: t.Type[EnumType], value: t.Union[int, EnumType]
|
||||
) -> t.Optional[EnumType]:
|
||||
if isinstance(value, int):
|
||||
return cls._create_pseudo_member_(value) # type: ignore
|
||||
return None # will raise the ValueError in Enum.__new__
|
||||
|
||||
@classmethod # type: ignore
|
||||
def _create_pseudo_member_(cls: t.Type[EnumType], value: int) -> EnumType:
|
||||
pseudo_member = cls._value2member_map_.get(value, None) # type: ignore
|
||||
if pseudo_member is None:
|
||||
new_member = int.__new__(cls, value) # type: ignore
|
||||
# I expect a name attribute to hold a string, hence str(value)
|
||||
# However, new_member._name_ = value works, too
|
||||
new_member._name_ = str(value)
|
||||
new_member._value_ = value
|
||||
pseudo_member = cls._value2member_map_.setdefault(value, new_member) # type: ignore
|
||||
return pseudo_member # type: ignore
|
||||
|
||||
# Monkey-patch the enum type with __missing__ method. So if a enum value
|
||||
# not found in the enum a new pseudo member is created.
|
||||
# The idea is taken from: https://stackoverflow.com/a/57179436
|
||||
enum_type._missing_ = _missing_ # type: ignore
|
||||
enum_type._create_pseudo_member_ = _create_pseudo_member_ # type: ignore
|
||||
|
||||
# save enum type
|
||||
self.enum_type = enum_type
|
||||
|
||||
# init adatper
|
||||
super(TEnum, self).__init__(subcon) # type: ignore
|
||||
|
||||
def _decode(self, obj: int, context: "cs.Context", path: "cs.PathType") -> EnumType:
|
||||
return self.enum_type(obj)
|
||||
|
||||
def _encode(
|
||||
self,
|
||||
obj: t.Union[int, str, EnumType],
|
||||
context: "cs.Context",
|
||||
path: "cs.PathType",
|
||||
) -> int:
|
||||
try:
|
||||
if isinstance(obj, str):
|
||||
return int(self.enum_type[obj])
|
||||
else:
|
||||
return int(self.enum_type(obj))
|
||||
except:
|
||||
raise cs.MappingError(
|
||||
"building failed, no mapping for %r" % (obj,), path=path
|
||||
)
|
||||
150
construct_typed/tstruct.py
Normal file
150
construct_typed/tstruct.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
import typing as t
|
||||
import textwrap
|
||||
import dataclasses
|
||||
from .generic_wrapper import *
|
||||
|
||||
|
||||
def StructField(
|
||||
subcon: Construct[ParsedType, BuildTypes],
|
||||
doc: t.Optional[str] = None,
|
||||
parsed: t.Optional[t.Callable[[t.Any, "cs.Context"], None]] = None,
|
||||
) -> ParsedType:
|
||||
"""
|
||||
Create a dataclass field for a "TStruct" and "TBitStruct" from a subcon.
|
||||
"""
|
||||
# Rename subcon, if doc or parsed are available
|
||||
if (doc is not None) or (parsed is not None):
|
||||
if doc is not None:
|
||||
doc = textwrap.dedent(doc)
|
||||
subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed)
|
||||
|
||||
if subcon.flagbuildnone is True:
|
||||
# some subcons have a predefined default value. all other have "None"
|
||||
default: t.Any = None
|
||||
if isinstance(subcon, (cs.Const, cs.Default)):
|
||||
if callable(subcon.value): # type: ignore
|
||||
raise ValueError("lamda as default is not supported")
|
||||
default = subcon.value # type: ignore
|
||||
|
||||
# if subcon builds from "None", set default to "None"
|
||||
field = dataclasses.field(
|
||||
default=default,
|
||||
init=False,
|
||||
metadata={"subcon": cs.Renamed(subcon, newdocs=doc)},
|
||||
)
|
||||
else:
|
||||
field = dataclasses.field(metadata={"subcon": subcon})
|
||||
|
||||
return field # type: ignore
|
||||
|
||||
|
||||
class _TStruct(Adapter[t.Any, t.Any, ParsedType, BuildTypes]):
|
||||
"""
|
||||
Base class for a typed struct, based on standard dataclasses.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, dataclass_type: t.Type[DataclassType], swapped: bool = False
|
||||
) -> None:
|
||||
if not dataclasses.is_dataclass(dataclass_type):
|
||||
raise TypeError(
|
||||
'The class "{}" is not a "dataclasses.dataclass"'.format(
|
||||
dataclass_type.__name__
|
||||
)
|
||||
)
|
||||
self.dataclass_type = dataclass_type
|
||||
self.swapped = swapped
|
||||
|
||||
# get all fields from the dataclass
|
||||
fields = dataclasses.fields(self.dataclass_type)
|
||||
if self.swapped:
|
||||
fields = tuple(reversed(fields))
|
||||
|
||||
# extract the construct formats from the struct_type
|
||||
subcon_fields = {}
|
||||
for field in fields:
|
||||
subcon_fields[field.name] = field.metadata["subcon"]
|
||||
|
||||
# init adatper
|
||||
super(_TStruct, self).__init__(self._create_subcon(subcon_fields)) # type: ignore
|
||||
|
||||
def _create_subcon(
|
||||
self, subcon_fields: t.Dict[str, t.Any]
|
||||
) -> Construct[t.Any, t.Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
def _decode(
|
||||
self, obj: "cs.Container[t.Any]", context: "cs.Context", path: "cs.PathType"
|
||||
) -> DataclassType:
|
||||
# get all fields from the dataclass
|
||||
fields = dataclasses.fields(self.dataclass_type)
|
||||
|
||||
# extract all fields from the container, that are used for create the dataclass object
|
||||
dc_init = {}
|
||||
for field in fields:
|
||||
if field.init:
|
||||
value = getattr(obj, field.name)
|
||||
dc_init[field.name] = value
|
||||
|
||||
# create object of dataclass
|
||||
dc = self.dataclass_type(**dc_init) # type: ignore
|
||||
|
||||
# extract all other values from the container, an pass it to the dataclass
|
||||
for field in fields:
|
||||
if not field.init:
|
||||
value = getattr(obj, field.name)
|
||||
setattr(dc, field.name, value)
|
||||
|
||||
return dc
|
||||
|
||||
def _encode(
|
||||
self, obj: DataclassType, context: "cs.Context", path: "cs.PathType"
|
||||
) -> t.Dict[str, t.Any]:
|
||||
# get all fields from the dataclass
|
||||
fields = dataclasses.fields(self.dataclass_type)
|
||||
|
||||
# extract all fields from the container, that are used for create the dataclass object
|
||||
ret_dict = {}
|
||||
for field in fields:
|
||||
value = getattr(obj, field.name)
|
||||
ret_dict[field.name] = value
|
||||
|
||||
return ret_dict
|
||||
|
||||
|
||||
class TStruct(_TStruct[ParsedType, BuildTypes]):
|
||||
"""
|
||||
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]):
|
||||
"""
|
||||
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]:
|
||||
return cs.BitStruct(**subcon_fields)
|
||||
44
construct_typed/tunion.py
Normal file
44
construct_typed/tunion.py
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
import enum
|
||||
import typing as t
|
||||
import textwrap
|
||||
import construct as cs
|
||||
import dataclasses
|
||||
from .generic_wrapper import *
|
||||
|
||||
|
||||
def UnionField(
|
||||
subcon: Construct[ParsedType, BuildTypes],
|
||||
doc: t.Optional[str] = None,
|
||||
parsed: t.Optional[t.Callable[[t.Any, "cs.Context"], None]] = None,
|
||||
) -> ParsedType:
|
||||
"""
|
||||
Create a dataclass field for a "TUnion" from a subcon.
|
||||
"""
|
||||
# Rename subcon, if doc or parsed are available
|
||||
if (doc is not None) or (parsed is not None):
|
||||
if doc is not None:
|
||||
doc = textwrap.dedent(doc)
|
||||
subcon = cs.Renamed(subcon, newdocs=doc, newparsed=parsed)
|
||||
|
||||
if subcon.flagbuildnone is True:
|
||||
# some subcons have a predefined default value. all other have "None"
|
||||
default: t.Any = None
|
||||
if isinstance(subcon, (cs.Const, cs.Default)):
|
||||
if callable(subcon.value): # type: ignore
|
||||
raise ValueError("lamda as default is not supported")
|
||||
default = subcon.value # type: ignore
|
||||
|
||||
# if subcon builds from "None", set default to "None"
|
||||
field = dataclasses.field(
|
||||
default=default,
|
||||
init=False,
|
||||
metadata={"subcon": cs.Renamed(subcon, newdocs=doc)},
|
||||
)
|
||||
else:
|
||||
field = dataclasses.field(metadata={"subcon": subcon})
|
||||
|
||||
return field # type: ignore
|
||||
|
||||
|
||||
class TUnion(Adapter[t.Any, t.Any, DataclassType, DataclassType]):
|
||||
pass # TODO
|
||||
1
setup.py
1
setup.py
|
|
@ -7,6 +7,7 @@ setup(
|
|||
packages=["construct-stubs", "construct_typed"],
|
||||
package_data={
|
||||
"construct-stubs": ["*.pyi", "lib/*.pyi"],
|
||||
"construct_typed": ["py.typed"],
|
||||
},
|
||||
include_package_data=True,
|
||||
license="MIT",
|
||||
|
|
|
|||
|
|
@ -4,132 +4,151 @@ import enum
|
|||
import dataclasses
|
||||
import typing as t
|
||||
from .declarativeunittest import common, raises
|
||||
from construct import (
|
||||
Int8ub,
|
||||
Int16ub,
|
||||
Const,
|
||||
Pass,
|
||||
Terminated,
|
||||
Padding,
|
||||
SizeofError,
|
||||
Byte,
|
||||
Bytes,
|
||||
Computed,
|
||||
this,
|
||||
)
|
||||
from construct_typed import TStruct, TBitStruct, TSubcon, TEnum
|
||||
import construct as cs
|
||||
import construct_typed as cst
|
||||
|
||||
|
||||
def test_typed_struct_1():
|
||||
def test_tstruct_1() -> None:
|
||||
@dataclasses.dataclass
|
||||
class TestDataclass:
|
||||
a: int = TSubcon(Int16ub)
|
||||
b: int = TSubcon(Int8ub)
|
||||
a: int = cst.StructField(cs.Int16ub)
|
||||
b: int = cst.StructField(cs.Int8ub)
|
||||
|
||||
common(cst.TStruct(TestDataclass), b"\x00\x01\x02", TestDataclass(a=1, b=2), 3)
|
||||
|
||||
|
||||
def test_tstruct_swapped() -> None:
|
||||
@dataclasses.dataclass
|
||||
class TestDataclass:
|
||||
a: int = cst.StructField(cs.Int16ub)
|
||||
b: int = cst.StructField(cs.Int8ub)
|
||||
|
||||
common(TStruct(TestDataclass), b"\x00\x01\x02", TestDataclass(a=1, b=2), 3)
|
||||
common(
|
||||
TStruct(TestDataclass, swapped=True),
|
||||
cst.TStruct(TestDataclass, swapped=True),
|
||||
b"\x02\x00\x01",
|
||||
TestDataclass(a=1, b=2),
|
||||
3,
|
||||
)
|
||||
normal = TStruct(TestDataclass)
|
||||
swapped = TStruct(TestDataclass, swapped=True)
|
||||
normal = cst.TStruct(TestDataclass)
|
||||
swapped = cst.TStruct(TestDataclass, swapped=True)
|
||||
assert str(normal.parse(b"\x00\x01\x02")) == str(swapped.parse(b"\x02\x00\x01"))
|
||||
|
||||
def test_typed_struct_2():
|
||||
|
||||
def test_tstruct_nested() -> None:
|
||||
@dataclasses.dataclass
|
||||
class TestDataclass:
|
||||
@dataclasses.dataclass
|
||||
class InnerDataclass:
|
||||
b: int = TSubcon(Byte)
|
||||
b: int = cst.StructField(cs.Byte)
|
||||
|
||||
a: InnerDataclass = TSubcon(TStruct(InnerDataclass))
|
||||
a: InnerDataclass = cst.StructField(cst.TStruct(InnerDataclass))
|
||||
|
||||
common(
|
||||
TStruct(TestDataclass),
|
||||
cst.TStruct(TestDataclass),
|
||||
b"\x01",
|
||||
TestDataclass(a=TestDataclass.InnerDataclass(b=1)),
|
||||
1,
|
||||
)
|
||||
|
||||
def test_typed_struct_3():
|
||||
|
||||
def test_tstruct_anonymus_fields_1() -> None:
|
||||
@dataclasses.dataclass
|
||||
class TestDataclass:
|
||||
_1: t.Optional[bytes] = TSubcon(Const(b"\x00"))
|
||||
_2: None = TSubcon(Padding(1))
|
||||
_3: None = TSubcon(Pass)
|
||||
_4: None = TSubcon(Terminated)
|
||||
_1: t.Optional[bytes] = cst.StructField(cs.Const(b"\x00"))
|
||||
_2: None = cst.StructField(cs.Padding(1))
|
||||
_3: None = cst.StructField(cs.Pass)
|
||||
_4: None = cst.StructField(cs.Terminated)
|
||||
|
||||
common(
|
||||
TStruct(TestDataclass),
|
||||
cst.TStruct(TestDataclass),
|
||||
bytes(2),
|
||||
TestDataclass(),
|
||||
SizeofError,
|
||||
cs.SizeofError,
|
||||
)
|
||||
|
||||
def test_typed_struct_4():
|
||||
|
||||
def test_tstruct_anonymus_fields_2() -> None:
|
||||
@dataclasses.dataclass
|
||||
class TestDataclass:
|
||||
_1: bytes = TSubcon(Bytes(this.missing))
|
||||
_1: int = cst.StructField(cs.Computed(7))
|
||||
_2: t.Optional[bytes] = cst.StructField(cs.Const(b"JPEG"))
|
||||
_3: None = cst.StructField(cs.Pass)
|
||||
_4: None = cst.StructField(cs.Terminated)
|
||||
|
||||
assert raises(TStruct(TestDataclass).sizeof) == SizeofError
|
||||
|
||||
def test_typed_struct_5():
|
||||
@dataclasses.dataclass
|
||||
class TestDataclass:
|
||||
_1: int = TSubcon(Computed(7))
|
||||
_2: t.Optional[bytes] = TSubcon(Const(b"JPEG"))
|
||||
_3: None = TSubcon(Pass)
|
||||
_4: None = TSubcon(Terminated)
|
||||
|
||||
d = TStruct(TestDataclass)
|
||||
d = cst.TStruct(TestDataclass)
|
||||
assert d.build(TestDataclass()) == d.build(TestDataclass())
|
||||
|
||||
|
||||
def test_typed_bit_struct():
|
||||
def test_tstruct_missing_dataclass() -> None:
|
||||
class TestDataclass:
|
||||
a: int = cst.StructField(cs.Int16ub)
|
||||
b: int = cst.StructField(cs.Int8ub)
|
||||
|
||||
assert raises(cst.TStruct, TestDataclass) == TypeError
|
||||
|
||||
|
||||
def test_tbitstruct() -> None:
|
||||
assert False
|
||||
|
||||
|
||||
def test_enum():
|
||||
def test_tenum() -> None:
|
||||
class E(enum.IntEnum):
|
||||
a = 1
|
||||
b = 2
|
||||
|
||||
a = TEnum(Byte, E)
|
||||
|
||||
common(TEnum(Byte, E), b"\x01", E.a, 1)
|
||||
common(TEnum(Byte, E), b"\x01", 1, 1)
|
||||
format = TEnum(Byte, E)
|
||||
common(cst.TEnum(cs.Byte, E), b"\x01", E.a, 1)
|
||||
common(cst.TEnum(cs.Byte, E), b"\x01", 1, 1)
|
||||
format = cst.TEnum(cs.Byte, E)
|
||||
obj = format.parse(b"\x01")
|
||||
assert obj == E.a
|
||||
assert obj == 1
|
||||
data = format.build("a")
|
||||
assert data == b"\x01"
|
||||
|
||||
common(TEnum(Byte, E), b"\x02", E.b, 1)
|
||||
common(TEnum(Byte, E), b"\x02", 2, 1)
|
||||
format = TEnum(Byte, E)
|
||||
common(cst.TEnum(cs.Byte, E), b"\x02", E.b, 1)
|
||||
common(cst.TEnum(cs.Byte, E), b"\x02", 2, 1)
|
||||
format = cst.TEnum(cs.Byte, E)
|
||||
obj = format.parse(b"\x02")
|
||||
assert obj == E.b
|
||||
assert obj == 2
|
||||
data = format.build("b")
|
||||
assert data == b"\x02"
|
||||
|
||||
common(TEnum(Byte, E), b"\x03", 3, 1)
|
||||
format = TEnum(Byte, E)
|
||||
|
||||
def test_tenum_missing_value() -> None:
|
||||
class E(enum.IntEnum):
|
||||
a = 1
|
||||
b = 2
|
||||
|
||||
common(cst.TEnum(cs.Byte, E), b"\x03", 3, 1)
|
||||
format = cst.TEnum(cs.Byte, E)
|
||||
obj = format.parse(b"\x03")
|
||||
assert int(obj) == 3
|
||||
data = format.build(3)
|
||||
assert data == b"\x03"
|
||||
|
||||
|
||||
def test_enum_in_struct():
|
||||
def test_tenum_no_int_enum() -> None:
|
||||
class E(enum.Enum):
|
||||
a = 1
|
||||
b = 2
|
||||
|
||||
assert raises(cst.TEnum, cs.Byte, E) == TypeError
|
||||
|
||||
|
||||
def test_tenum_in_tstruct() -> None:
|
||||
class TestEnum(enum.IntEnum):
|
||||
a = 1
|
||||
b = 2
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TestDataclass:
|
||||
a: TestEnum = TSubcon(TEnum(Int8ub, TestEnum))
|
||||
b: int = TSubcon(Int8ub)
|
||||
a: TestEnum = cst.StructField(cst.TEnum(cs.Int8ub, TestEnum))
|
||||
b: int = cst.StructField(cs.Int8ub)
|
||||
|
||||
common(TStruct(TestDataclass), b"\x00\x01\x02", TestDataclass(a=TestEnum.a, b=TestEnum.b), 3)
|
||||
common(
|
||||
cst.TStruct(TestDataclass),
|
||||
b"\x00\x01\x02",
|
||||
TestDataclass(a=TestEnum.a, b=TestEnum.b),
|
||||
3,
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue