diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..c695aca --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,10 @@ +{ + "python.languageServer": "Pylance", + "python.testing.unittestEnabled": false, + "python.testing.nosetestsEnabled": false, + "python.testing.pytestEnabled": true, + "pythonTestExplorer.testFramework": "pytest", + "python.formatting.provider": "black", + "python.analysis.typeCheckingMode": "strict", + "python.analysis.autoImportCompletions": false +} \ No newline at end of file diff --git a/README.md b/README.md index 1db6584..20d24d0 100644 --- a/README.md +++ b/README.md @@ -1 +1,33 @@ -# construct-typing \ No newline at end of file +# construct-typing + +## +This project is an extension of the python package `construct`. This Repository consitst of two packages: + +- **construct-stubs**: Adding .pyi for the whole `construct` package (according to [PEP 561 stub-only packages](https://www.python.org/dev/peps/pep-0561/#stub-only-packages)) +- **construct-typed**: Adding the additional classes `TypedStruct` and `TypedEnum` that help with autocompletion in cooperation with the stubs. + +## Motivation + + +## Examples + +An example of the added `TypedStruct` class: + +```python +from construct import * +from construct_typed import * + +class Image(TypedContainer): + signature: Subcon(Const(b"BMP")) + width: Subcon(Int8ub()) + height: Subcon(Int8ub()) + pixels: Subcon(Array(cs.this.width * cs.this.height, Byte())) + +format = TypedStruct(Image) +obj = Image(width=3, height=2, pixels=[7, 8, 9, 11, 12, 13]) +print(format.build(obj)) +print(format.parse(b"BMP\x03\x02\x07\x08\t\x0b\x0c\r")) +``` + +An example of the added `TypedEnum` class: + diff --git a/construct-stubs/__init__.pyi b/construct-stubs/__init__.pyi new file mode 100644 index 0000000..1236325 --- /dev/null +++ b/construct-stubs/__init__.pyi @@ -0,0 +1,181 @@ +from construct.core import * +from construct.expr import * +from construct.debug import * +from construct.version import * +from construct import lib + + +#=============================================================================== +# exposed names +#=============================================================================== +__all__ = [ + '__author__', + '__version__', + 'abs_', + 'AdaptationError', + 'Adapter', + 'Aligned', + 'AlignedStruct', + 'Array', + 'Bit', + 'BitsInteger', + 'BitsSwapped', + 'BitStruct', + 'BitwisableString', + 'Bitwise', + 'Byte', + 'Bytes', + 'BytesInteger', + 'ByteSwapped', + 'Bytewise', + 'CancelParsing', + 'Check', + 'CheckError', + 'Checksum', + 'ChecksumError', + 'Compiled', + 'Compressed', + 'Computed', + 'Const', + 'ConstError', + 'Construct', + 'ConstructError', + 'Container', + 'CString', + 'Debugger', + 'Default', + 'Double', + 'Enum', + 'EnumInteger', + 'EnumIntegerString', + 'Error', + 'ExplicitError', + 'ExprAdapter', + 'ExprSymmetricAdapter', + 'ExprValidator', + 'Filter', + 'FixedSized', + 'Flag', + 'FlagsEnum', + 'FocusedSeq', + 'FormatField', + 'FormatFieldError', + 'FuncPath', + 'globalPrintFalseFlags', + 'globalPrintFullStrings', + 'GreedyBytes', + 'GreedyRange', + 'GreedyString', + 'Half', + 'Hex', + 'HexDump', + 'If', + 'IfThenElse', + 'Index', + 'IndexFieldError', + 'Indexing', + 'Int', + 'IntegerError', + 'Lazy', + 'LazyArray', + 'LazyBound', + 'LazyContainer', + 'LazyListContainer', + 'LazyStruct', + 'len_', + 'lib', + 'list_', + 'ListContainer', + 'Long', + 'Mapping', + 'MappingError', + 'max_', + 'min_', + 'NamedTuple', + 'NamedTupleError', + 'Nibble', + 'NoneOf', + 'NullStripped', + 'NullTerminated', + 'Numpy', + 'obj_', + 'Octet', + 'OneOf', + 'Optional', + 'Padded', + 'PaddedString', + 'Padding', + 'PaddingError', + 'PascalString', + 'Pass', + 'Path', + 'Path2', + 'Peek', + 'Pickled', + 'Pointer', + 'possiblestringencodings', + 'Prefixed', + 'PrefixedArray', + 'Probe', + 'ProcessRotateLeft', + 'ProcessXor', + 'RangeError', + 'RawCopy', + 'Rebuffered', + 'RebufferedBytesIO', + 'Rebuild', + 'release_date', + 'Renamed', + 'RepeatError', + 'RepeatUntil', + 'RestreamData', + 'Restreamed', + 'RestreamedBytesIO', + 'RotationError', + 'Seek', + 'Select', + 'SelectError', + 'Sequence', + 'setGlobalPrintFalseFlags', + 'setGlobalPrintFullStrings', + 'setGlobalPrintPrivateEntries', + 'Short', + 'Single', + 'SizeofError', + 'Slicing', + 'StopFieldError', + 'StopIf', + 'stream_iseof', + 'stream_read', + 'stream_read_entire', + 'stream_seek', + 'stream_size', + 'stream_tell', + 'stream_write', + 'StreamError', + 'StringEncoded', + 'StringError', + 'Struct', + 'Subconstruct', + 'sum_', + 'Switch', + 'SwitchError', + 'SymmetricAdapter', + 'Tell', + 'Terminated', + 'TerminatedError', + 'this', + 'Timestamp', + 'TimestampError', + 'Transformed', + 'Tunnel', + 'Union', + 'UnionError', + 'ValidationError', + 'Validator', + 'VarInt', + 'version', + 'version_string', +] +__all__ += ["Int%s%s%s" % (n,us,bln) for n in (8,16,24,32,64) for us in "us" for bln in "bln"] +__all__ += ["Float%s%s" % (n,bln) for n in (16,32,64) for bln in "bln"] diff --git a/construct-stubs/debug.pyi b/construct-stubs/debug.pyi new file mode 100644 index 0000000..e69de29 diff --git a/construct-stubs/lib/__init__.pyi b/construct-stubs/lib/__init__.pyi new file mode 100644 index 0000000..46658a7 --- /dev/null +++ b/construct-stubs/lib/__init__.pyi @@ -0,0 +1,53 @@ +from construct.lib.containers import * +from construct.lib.binary import * +from construct.lib.bitstream import * +from construct.lib.hex import * +from construct.lib.py3compat import * + +__all__ = [ + 'bits2bytes', + 'bits2integer', + 'byte2int', + 'bytes2bits', + 'bytes2integer', + 'bytes2integers', + 'bytes2str', + 'bytestringtype', + 'Container', + 'globalPrintFalseFlags', + 'globalPrintFullStrings', + 'HexDisplayedBytes', + 'HexDisplayedDict', + 'HexDisplayedInteger', + 'hexdump', + 'HexDumpDisplayedBytes', + 'HexDumpDisplayedDict', + 'hexlify', + 'hexundump', + 'int2byte', + 'integer2bits', + 'integer2bytes', + 'integers2bytes', + 'integertypes', + 'iteratebytes', + 'iterateints', + 'ListContainer', + 'PY', + 'PY2', + 'PY3', + 'PYPY', + 'RebufferedBytesIO', + 'reprstring', + 'RestreamedBytesIO', + 'setGlobalPrintFalseFlags', + 'setGlobalPrintFullStrings', + 'setGlobalPrintPrivateEntries', + 'str2bytes', + 'stringtypes', + 'swapbitsinbytes', + 'swapbytes', + 'swapbytesinbits', + 'trimstring', + 'unhexlify', + 'unicodestringtype', +] diff --git a/construct-stubs/lib/binary.pyi b/construct-stubs/lib/binary.pyi new file mode 100644 index 0000000..e69de29 diff --git a/construct-stubs/lib/bitstream.pyi b/construct-stubs/lib/bitstream.pyi new file mode 100644 index 0000000..e69de29 diff --git a/construct-stubs/lib/py3compat.pyi b/construct-stubs/lib/py3compat.pyi new file mode 100644 index 0000000..e69de29 diff --git a/construct-stubs/version.pyi b/construct-stubs/version.pyi new file mode 100644 index 0000000..5dbe453 --- /dev/null +++ b/construct-stubs/version.pyi @@ -0,0 +1,4 @@ +from typing import Tuple +version: Tuple[int, int, int] +version_string: str +release_date: str diff --git a/construct_typed/__init__.py b/construct_typed/__init__.py new file mode 100644 index 0000000..d525e69 --- /dev/null +++ b/construct_typed/__init__.py @@ -0,0 +1,121 @@ +from enum import IntEnum +from typing import Any, Type, Dict, TYPE_CHECKING, TypeVar, Union +import typing +from construct.core import Construct, Adapter +from construct.lib.containers import Container + + +#=============================================================================== +# mappings +#=============================================================================== +if TYPE_CHECKING: + EnumType = TypeVar("EnumType", bound=IntEnum) + + class TypedEnum(Adapter[EnumType, Union[int, str, IntEnum], int, int]): + def __init__(self, subcon: Construct[int, int], enum_type: Type[EnumType]): ... + +# def wrap_enum(enum_type: Type[EnumType]) -> +else: + class TypedEnum(Adapter): + def __init__(self, subcon, enum_type): + super(TypedEnum, self).__init__(subcon) + + @classmethod + def _missing_(cls, value): + 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, value): + 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_ + + self.enum_type = enum_type + + def _decode(self, obj, context, path): + if isinstance(obj, str): + return int(self.enum_type[obj]) + else: + return self.enum_type(obj) + + def _encode(self, obj, context, path): + try: + if isinstance(obj, str): + return int(self.enum_type[obj]) + else: + return int(self.enum_type(obj)) + except: + raise MappingError("building failed, no mapping for %r" % (obj,), path=path) + + + + +# =============================================================================== +# structures and sequences +# =============================================================================== +ParsedType = TypeVar("ParsedType") +BuildTypes = TypeVar("BuildTypes") + +# In an optimal way the code look like this (with python 3.9): +# +# def Subcon(subcon: Construct[ParsedType, BuildTypes]) -> Type[ParsedType]: +# return Annotated[ParsedType, subcon] +# +# But this only works if the type annotations are directly in the source +# code. However, in this case the type annotations are stored in a separate +# stub (.pyi) file, which is not available at runtime. +# +# Therefore a small hack is used here. During the type checking, the type +# of the subcon is determined with the help of the stub files and returned. +# At runtime the actual subcon and not its type is returned, so that the +# "TypedStruct" can use this for creating the struct. +# +# This has the disadvantage that it can cause problems if the type is evaluated +# at runtime... +if TYPE_CHECKING: + def Subcon(subcon: Construct[ParsedType, BuildTypes]) -> Type[ParsedType]: ... +else: + def Subcon(subcon: Construct) -> Construct: + return subcon + +if TYPE_CHECKING: + class TypedContainer(Container[Any]): + ... +else: + class TypedContainer(Container): + pass + +if TYPE_CHECKING: + ContainerType = TypeVar("ContainerType", bound=TypedContainer) + class TypedStruct(Construct[ContainerType, Dict[str, Any]]): + def __init__(self, container_type: Type[ContainerType], swapped: bool = False) -> None: ... +else: + class TypedStruct(Struct): + def __init__(self, container_type, swapped = False): + if not issubclass(container_type, TypedContainer): + raise TypeError("the subcon has to be a TypedContainer") + + # extract the construct formats from the struct_type + subcons = {} + subcon_formats = typing.get_type_hints(container_type) + subcon_items = subcon_formats.items() + if swapped: + subcon_items = reversed(subcon_items) + for subcon_name, subcon_format in subcon_items: + subcons[subcon_name] = subcon_format + + super(TypedStruct, self).__init__(**subcons) diff --git a/construct_typed/as.py b/construct_typed/as.py new file mode 100644 index 0000000..18f2328 --- /dev/null +++ b/construct_typed/as.py @@ -0,0 +1,15 @@ +from construct.core import * +from construct_typed import * + +a = Const(b"asd") + +class Image(TypedContainer): + signature: Subcon(Const(b"asd")) + width: Subcon(Int8ub) + height: Subcon(Int8ub) + pixels: Subcon(Array(cs.this.width * cs.this.height, Byte())) + +format = TypedStruct(Image) +obj = Image(width=3, height=2, pixels=[7, 8, 9, 11, 12, 13]) +print(format.build(obj)) +print(format.parse(b"BMP\x03\x02\x07\x08\t\x0b\x0c\r")) \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..93b6242 --- /dev/null +++ b/setup.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python +from setuptools import setup + +setup( + name="construct-typing", + version="0.0.1", + packages=["construct-stubs", "construct_typed"], + package_data={ + "construct-stubs": ["*.pyi", "lib/*.pyi"], + }, + include_package_data=True, + license="MIT", + description="Extension for the python package 'construct' that adds typing features", + long_description=open("README.md").read(), + platforms=["Windows"], + url="http://construct.readthedocs.org", + author="Tim Riddermann", + python_requires=">=3.6", + install_requires=["construct==2.10.56"], + keywords=[ + "construct", + "kaitai", + "declarative", + "data structure", + "struct", + "binary", + "symmetric", + "parser", + "builder", + "parsing", + "building", + "pack", + "unpack", + "packer", + "unpacker", + "bitstring", + "bytestring", + "annotation", + "type hint", + "typing", + "typed", + "bitstruct", + "PEP 561" + ], + classifiers=[ + "Development Status :: 3 - Alpha", + "License :: OSI Approved :: MIT License", + "Intended Audience :: Developers", + "Topic :: Software Development :: Libraries :: Python Modules", + "Topic :: Software Development :: Build Tools", + "Topic :: Software Development :: Code Generators", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.6", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: Implementation :: CPython", + "Programming Language :: Python :: Implementation :: PyPy", + "Typing :: Typed" + ], +) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..f4df255 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,2 @@ +import sys +sys.path.insert(0, 'tests') \ No newline at end of file diff --git a/tests/declarativeunittest.py b/tests/declarativeunittest.py new file mode 100644 index 0000000..adc8115 --- /dev/null +++ b/tests/declarativeunittest.py @@ -0,0 +1,68 @@ +import pytest +xfail = pytest.mark.xfail +skip = pytest.mark.skip +skipif = pytest.mark.skipif + +import os, math, random, collections, itertools, io, hashlib, binascii + +from construct import * +from construct.lib import * + +class ZeroIO(io.BufferedIOBase): + def read(self, __size = None): + if __size is not None: + return bytes(__size) + else: + return bytes(0) + + def read1(self, __size = 0): + return bytes(__size) + +ident = lambda x: x +devzero = ZeroIO() + + +def raises(func, *args, **kw): + try: + return func(*args, **kw) + except Exception as e: + return e.__class__ + + +def common(format, datasample, objsample, sizesample=SizeofError, **kw): + obj = format.parse(datasample, **kw) + assert obj == objsample + data = format.build(objsample, **kw) + assert data == datasample + # following are implied by above (re-parse and re-build) + # assert format.parse(format.build(obj)) == obj + # assert format.build(format.parse(data)) == data + if isinstance(sizesample, int): + size = format.sizeof(**kw) + assert size == sizesample + else: + size = raises(format.sizeof, **kw) + assert size == sizesample + + +def commonhex(format, hexdata): + commonbytes(format, binascii.unhexlify(hexdata)) + + +def commondumpdeprecated(format, filename): + filename = "tests/deprecated_gallery/blobs/" + filename + with open(filename,'rb') as f: + data = f.read() + commonbytes(format, data) + + +def commondump(format, filename): + filename = "tests/gallery/blobs/" + filename + with open(filename,'rb') as f: + data = f.read() + commonbytes(format, data) + + +def commonbytes(format, data): + obj = format.parse(data) + data2 = format.build(obj) diff --git a/tests/declarativeunittest.pyi b/tests/declarativeunittest.pyi new file mode 100644 index 0000000..d0f69b1 --- /dev/null +++ b/tests/declarativeunittest.pyi @@ -0,0 +1,32 @@ +from typing import Callable, TypeVar, Union, Any, Type, BinaryIO +from construct import * +from construct.lib import * + +Buffer = Union[bytes, memoryview, bytearray] +ParsedType = TypeVar("ParsedType") +BuildTypes = TypeVar("BuildTypes") + +devzero: BinaryIO + +def raises( + func: Callable[..., Any], *args: Any, **kw: Any +) -> Union[Any, Exception]: ... +def common( + format: Construct[ParsedType, BuildTypes], + datasample: Buffer, + objsample: BuildTypes, + sizesample: Union[int, Type[Exception]] = ..., + **kw: Any +) -> None: ... +def commonhex( + format: Construct[ParsedType, BuildTypes], hexdata: str +) -> None: ... +def commondumpdeprecated( + format: Construct[ParsedType, BuildTypes], filename: str +) -> None: ... +def commondump( + format: Construct[ParsedType, BuildTypes], filename: str +) -> None: ... +def commonbytes( + format: Construct[ParsedType, BuildTypes], data: ParsedType +) -> None: ... diff --git a/tests/test_typed.py b/tests/test_typed.py new file mode 100644 index 0000000..c878caf --- /dev/null +++ b/tests/test_typed.py @@ -0,0 +1,110 @@ +# -*- coding: utf-8 -*- + +import enum +from .declarativeunittest import common, raises +from construct.core import ( + Int8ub, + Int16ub, + Const, + Pass, + Terminated, + Padding, + SizeofError, + Byte, + Bytes, + Computed, + this, +) +from construct.lib import Container +from construct_typed import TypedContainer, TypedStruct, TypedEnum, Subcon + + +def test_typed_struct(): + class Container1(TypedContainer): + a: Subcon(Int16ub) + b: Subcon(Int8ub) + + common(TypedStruct(Container1), b"\x00\x01\x02", Container1(a=1, b=2), 3) + + class Container1Reversing(TypedContainer): + a: Subcon(Int16ub) + b: Subcon(Int8ub) + + common( + TypedStruct(Container1Reversing, swapped=True), + b"\x02\x00\x01", + Container(a=1, b=2), + 3, + ) + + class Container2(TypedContainer): + class InnerContainer(TypedContainer): + b: Subcon(Byte) + + a: Subcon(TypedStruct(InnerContainer)) + + common(TypedStruct(Container2), b"\x01", Container(a=Container(b=1)), 1) + + # TODO: How to get anonymus subcons? + class Container3(TypedContainer): + anonymus1: Subcon(Const(b"\x00")) + anonymus2: Subcon(Padding(1)) + anonymus3: Subcon(Pass) + anonymus4: Subcon(Terminated) + + common( + TypedStruct(Container3), + bytes(2), + dict(anonymus1=b"\x00", anonymus2=None, anonymus3=None, anonymus4=None), + SizeofError, + ) + + class Container4(TypedContainer): + missingkey: Subcon(Byte) + + assert raises(TypedStruct(Container4).build, {}) == KeyError + + # TODO: How to get anonymus subcons? + class Container5(TypedContainer): + anonymus1: Subcon(Bytes(this.missing)) + + assert raises(TypedStruct(Container5).sizeof) == SizeofError + + # TODO: How to get anonymus subcons? + class Container6(TypedContainer): + anonymus1: Subcon(Computed(7)) + anonymus2: Subcon(Const(b"JPEG")) + anonymus3: Subcon(Pass) + anonymus4: Subcon(Terminated) + + d = TypedStruct(Container6) + assert d.build({}) == d.build({}) + + +def test_enum(): + class E(enum.IntEnum): + a = 1 + b = 2 + + common(TypedEnum(Byte, E), b"\x01", E.a, 1) + common(TypedEnum(Byte, E), b"\x01", 1, 1) + format = TypedEnum(Byte, E) + obj = format.parse(b"\x01") + assert obj == E.a + data = format.build("a") + assert data == b"\x01" + + common(TypedEnum(Byte, E), b"\x02", E.b, 1) + common(TypedEnum(Byte, E), b"\x02", 2, 1) + format = TypedEnum(Byte, E) + obj = format.parse(b"\x02") + assert obj == E.b + data = format.build("b") + assert data == b"\x02" + + common(TypedEnum(Byte, E), b"\x03", 3, 1) + format = TypedEnum(Byte, E) + obj = format.parse(b"\x03") + assert int(obj) == 3 + data = format.build(3) + assert data == b"\x03"