sx/sx/types/seq.py

from typing import Optional as O, Union as U, Callable, Any, List, Sequence, Mapping, Generic as G, TypeVar, Tuple as Tu
from types import FunctionType

from ..core.base import PossibleDynamic as D, Type, Context, PathElement
from ..core.io import Stream, add_sizes
from ..core import to_type


T = TypeVar('T', bound=Type)

class Arr(G[T], Type[List[T]]):
    def __init__(self, child: D[T], count: O[D[int]] = None, stop: O[U[D[Any], Callable[[Any], bool]]] = None, include_stop: D[bool] = False) -> None:
        self.child = child
        self.count = count
        self.stop = stop
        self.include_stop = include_stop

    def parse(self, context: Context, stream: Stream) -> List[T]:
        child = context.get(self.child)
        count = context.get(self.count)
        stop = context.get(self.stop)
        include_stop = context.get(self.include_stop)

        value = []
        while True:
            i = len(value)
            if count is not None and i >= count:
                break
            c = to_type(child)
            with context.enter(i, c):
                try:
                    elem = context.parse(c, stream)
                except EOFError:
                    if count is None:
                        break
                    raise
            if stop is not None:
                if isinstance(stop, FunctionType):
                    should_stop = stop(elem)
                else:
                    should_stop = elem == stop
                if should_stop:
                    if include_stop:
                        value.append(elem)
                    break
            value.append(elem)
        
        return value

    def dump(self, context: Context, stream: Stream, value: List[T]) -> None:
        child = context.get(self.child)
        count = context.get(self.count)
        stop = context.get(self.stop)
        include_stop = context.get(self.include_stop)

        if stop is not None and not isinstance(stop, FunctionType) and not include_stop:
            value += [stop]

        for i, elem in enumerate(value):
            c = to_type(child)
            with context.enter(i, c):
                context.dump(child, stream, elem)

        context.put(self.count, len(value))

    def get_sizes(self, context: Context, value: O[List[T]], n: int) -> Mapping[str, int]:
        child = context.peek(self.child)
        stop = context.peek(self.stop)

        sizes = []
        for i in range(n):
            c = to_type(child)
            if value is not None:
                elem = value[i]
            else:
                elem = None
            with context.enter(i, c):
                size = context.sizeof(c, elem)
            sizes.append(size)

        if stop is not None and not isinstance(stop, FunctionType):
            sizes.append(context.sizeof(child, stop))

        return sizes

    def sizeof(self, context: Context, value: O[List[T]]) -> O[Mapping[str, int]]:
        if value is not None:
            count = len(value)
        else:
            count = context.peek(self.count)
        if count is None:
            return None
        sizes = self.get_sizes(context, value, count)
        return add_sizes(*sizes) if sizes else 0

    def offsetof(self, context: Context, path: Sequence[PathElement], value: O[List[T]]) -> O[int]:
        if not path:
            return 0

        i = path[0]
        path = path[1:]
        if not isinstance(i, int):
            raise ValueError('path element for array must be integer')

        child = context.peek(self.child)
        sizes = self.get_sizes(context, value, i)
        if path:
            c = to_type(child)
            with context.enter(i, c):
                sizes.append(context.offsetof(c, path, value[i] if value is not None else None))
        return add_sizes(*sizes) if sizes else 0       

    def default(self, context: Context) -> List[T]:
        child = context.peek(self.child)
        count = context.peek(self.count)
        if count is None or child is None:
            value = []
        else:
            value = [context.default(child) for _ in range(count)]
        return value

    def __str__(self) -> str:
        if self.count is not None:
            count = repr(self.count)
        else:
            count = ''
        return f'{self.child}[{count}]'

    def __repr__(self) -> str:
        return f'<{__name__}.{self.__class__.__name__}({self.child!r}, count: {self.count!r}, stop: {self.stop!r}, include_stop: {self.include_stop!r})>'


class Tuple(Type):
    def __init__(self, *children: Type) -> None:
        self.children = children

    def parse(self, context: Context, stream: Stream) -> Tu:
        values = []

        for i, child in enumerate(self.children):
            c = to_type(child)
            with context.enter(i, c):
                elem = context.parse(c, stream)
            values.append(elem)
        
        return tuple(values)

    def dump(self, context: Context, stream: Stream, value: Tu) -> None:
        for i, (child, elem) in enumerate(zip(self.children, value)):
            c = to_type(child)
            with context.enter(i, c):
                context.dump(c, stream, elem)

    def sizeof(self, context: Context, value: O[Tu]) -> O[int]:
        sizes = []

        if value is None:
            value = [None] * len(self.children)
        for i, (child, elem) in enumerate(zip(self.children, value)):
            c = to_type(child)
            with context.enter(i, c):
                sizes.append(context.sizeof(c, elem))
        
        return add_sizes(*sizes)

    def offsetof(self, context: Context, path: Sequence[PathElement], value: O[Tu]) -> O[int]:
        if not path:
            return 0

        n = path[0]
        path = path[1:]
        if not isinstance(n, int):
            raise ValueError('path element for tuple must be integer')

        sizes = []

        if value is None:
            value = [None] * len(self.children)
        for i, (child, elem) in enumerate(zip(self.children, value)):
            if i >= n:
                break
            c = to_type(child)
            with context.enter(i, c):
                sizes.append(context.sizeof(c, elem))
        
        if path:
            c = to_type(child)
            with context.enter(n, c):
                sizes.append(context.offsetof(c, path, elem))

        return add_sizes(*sizes)

    def default(self, context: Context) -> Tu:
        return tuple(context.default(c) for c in self.children)

    def __str__(self) -> str:
        return '(' + ', '.join(str(c) for c in self.children) + ')'

    def __repr__(self) -> str:
        return f'<{__name__}.{self.__class__.__name__}(' + ', '.join(repr(c) for c in self.children) + ')>'