io: add AlignTo and AlignedTo types

This commit is contained in:
Shiz 2021-06-26 03:36:53 +02:00
parent c4ddf590a5
commit 7a7a2a90cc
3 changed files with 90 additions and 0 deletions

View File

@ -9,11 +9,13 @@ from .types.num import *
from .types.struct import StructType, Struct
from .types.seq import Arr, Tuple
from .types.transforms import Default, Sized, Ref, Transform, Mapped
from .types.io import AlignTo, AlignedTo
__all__ = [x.__name__ for x in {
parse, dump, sizeof, offsetof, default,
Context, Type, Stream, Segment,
Wrapper, Default, Sized, Ref, Transform, Mapped,
AlignTo, AlignedTo,
Data,
Int, Bool, Float,
Arr, Tuple,

View File

@ -25,6 +25,12 @@ def seeking(fd: BinaryIO, pos: Pos, whence: int = os.SEEK_SET) -> Generator[Bina
fd.seek(oldpos, os.SEEK_SET)
def stretch(b: bytes, count: int) -> bytes:
b *= count // len(b)
b += b[:count - len(b)]
return b
def indent(s: str, count: int, start: bool = False) -> str:
""" Indent all lines of a string. """
lines = s.splitlines()

82
sx/types/io.py Normal file
View File

@ -0,0 +1,82 @@
from typing import Generic as G, TypeVar, Union as U, Optional as O, Sequence
import os
from ..core.base import Type, Context, PossibleDynamic, PathElement
from ..core.io import Stream, Pos
from ..core.meta import Wrapper
from ..core.util import stretch
T = TypeVar('T')
class AlignTo(G[T], Wrapper[T]):
__slots__ = ('alignment', 'value')
def __init__(self, child: Type[T], alignment: U[PossibleDynamic, Pos], value: U[PossibleDynamic, bytes] = b'\x00') -> None:
super().__init__(child)
self.alignment = alignment
self.value = value
def parse(self, context: Context, stream: Stream) -> T:
value = super().parse(context, stream)
align = context.get(self.alignment)
adjustment = stream.tell() % align
if adjustment:
stream.seek(align - adjustment, os.SEEK_CUR)
return value
def dump(self, context: Context, stream: Stream, value: T) -> None:
super().dump(context, stream, value)
align = context.get(self.alignment)
adjustment = stream.tell() % align
padding = stretch(context.get(self.value), align - adjustment)
stream.write(padding)
def sizeof(self, context: Context, value: O[T]) -> O[Pos]:
# TODO
return None
def __str__(self) -> str:
return f'{super().__str__()}%{self.align}'
def __repr__(self) -> str:
return f'{__name__}.AlignTo({super().__repr__()}, alignment={self.alignment!r}, value={self.value!r})'
class AlignedTo(G[T], Wrapper[T]):
__slots__ = ('alignment', 'value')
def __init__(self, child: Type[T], alignment: U[PossibleDynamic, Pos], value: U[PossibleDynamic, bytes] = b'\x00') -> None:
super().__init__(child)
self.alignment = alignment
self.value = value
def parse(self, context: Context, stream: Stream) -> T:
align = context.get(self.alignment)
adjustment = stream.tell() % align
if adjustment:
stream.seek(align - adjustment, os.SEEK_CUR)
return super().parse(context, stream)
def dump(self, context: Context, stream: Stream, value: T) -> None:
align = context.get(self.alignment)
adjustment = stream.tell() % align
padding = stretch(context.get(self.value), align - adjustment)
stream.write(padding)
super().dump(context, stream, value)
def sizeof(self, context: Context, value: O[T]) -> O[Pos]:
# TODO
return None
def offsetof(self, context: Context, path: Sequence[PathElement], value: O[T]) -> O[Pos]:
# TODO
return None
def __str__(self) -> str:
return f'{super().__str__()}%{self.align}'
def __repr__(self) -> str:
return f'{__name__}.AlignTo({super().__repr__()}, alignment={self.alignment!r}, value={self.value!r})'