sx/sx/core/meta.py

from typing import List, Optional as O, Generic as G, Sequence, Tuple, TypeVar, Any

import os
from .base import Type, Context, PathElement, Error, to_type
from .io import Stream, Segment, Pos
from .expr import BaseExpr


T = TypeVar('T')

class Wrapper(G[T], Type[T]):
    def __init__(self, child: Type[T]) -> None:
        self.child = child

    def parse(self, context: Context, stream: Stream) -> T:
        child = to_type(self.child)
        with context.enter(None, child):
            return context.parse(child, stream)

    def dump(self, context: Context, stream: Stream, value: O[T]) -> None:
        child = to_type(self.child)
        with context.enter(None, child):
            context.dump(child, stream, value)

    def sizeof(self, context: Context, value: O[T]) -> O[int]:
        child = to_type(self.child)
        with context.enter(None, child):
            return context.sizeof(child, value)

    def offsetof(self, context: Context, path: Sequence[PathElement], value: O[T]) -> O[int]:
        child = to_type(self.child)
        with context.enter(None, child):
            return context.offsetof(child, path, value)

    def default(self, context: Context) -> T:
        child = to_type(self.child)
        with context.enter(None, child):
            return context.default(child)

    def __str__(self) -> str:
        return str(to_type(self.child))

    def __repr__(self) -> str:
        return repr(to_type(self.child))

class Generic(G[T], Type[T], BaseExpr[T]):
    __slots__ = ('name', 'stack')

    def __init__(self, name: str, stack: O[List[Any]] = None) -> None:
        self.name = name
        self.stack = stack or []

    def push(self, value: Any) -> None:
        if isinstance(value, Generic):
            self.stack.append(value.stack[-1])
        else:
            self.stack.append(value)

    def pop(self) -> None:
        self.stack.pop()

    def _sx_get_(self, pop: bool = False) -> T:
        return self.stack[-1]

    def _sx_peek_(self, pop: bool = False) -> T:
        return self.stack[-1]

    def _sx_is_const_(self) -> bool:
        return False

    def _get_sx_type_(self, ident: Any) -> Type:
        return to_type(self.stack[-1])

    def parse(self, context: Context, stream: Stream) -> T:
        if not self.stack:
            raise Error(context, 'unresolved generic')
        child = to_type(self.stack[-1])
        with context.enter(None, child):
            return context.parse(child, stream)

    def dump(self, context: Context, stream: Stream, value: T) -> None:
        if not self.stack:
            raise Error(context, 'unresolved generic')
        child = to_type(self.stack[-1])
        with context.enter(None, child):
            context.dump(child, stream, value)

    def sizeof(self, context: Context, value: O[T]) -> O[int]:
        if not self.stack:
            return None
        child = to_type(self.stack[-1])
        with context.enter(None, child):
            return context.sizeof(child, value)

    def offsetof(self, context: Context, path: Sequence[PathElement], value: O[T]) -> O[int]:
        if not self.stack:
            return None
        child = to_type(self.stack[-1])
        with context.enter(None, child):
            return context.offsetof(child, path, value)

    def default(self, context: Context) -> T:
        if not self.stack:
            raise Error(context, 'unresolved generic')
        child = to_type(self.stack[-1])
        with context.enter(None, child):
            return context.default(child)

    def __str__(self) -> str:
        if self.stack:
            return f'${self.name}:{to_type(self.stack[-1])}'
        return f'${self.name}:unresolved'

    def __repr__(self) -> str:
        return f'{__name__}.Generic({self.name!r}, {self.stack!r})'

    def __deepcopy__(self, memo: Any) -> Any:
        return self

class TypeSource(G[T], Wrapper[T], BaseExpr[T]):
    def __init__(self, child: Type[T], count: int) -> None:
        super().__init__(child)
        self.stack: list[Tuple[Context, Segment, Stream, Pos, T]] = []
        self.pstack: list[T] = []
        self.count = count

    def parse(self, context: Context, stream: Stream) -> T:
        pos = stream.tell()
        value = super().parse(context, stream)
        for _ in range(self.count):
            self.stack.append((context, context.segment, stream, pos, value))
        return value

    def dump(self, context: Context, stream: Stream, value: T) -> None:
        pos = stream.tell()
        for _ in range(self.count):
            self.stack.append((context, context.segment, stream, pos, value))
        super().dump(context, stream, value)

    def sizeof(self, context: Context, value: O[T]) -> None:
        for _ in range(self.count):
            self.pstack.append(value)
        return super().sizeof(context, value)

    def offsetof(self, context: Context, path: Sequence[PathElement], value: O[T]) -> None:
        for _ in range(self.count):
            self.pstack.append(value)
        return super().offsetof(context, path, value)

    def default(self, context: Context) -> T:
        value = super().default(context)
        for _ in range(self.count):
            self.pstack.append(value)
        return value

    def _sx_get_(self, pop: bool = True) -> T:
        _, _, _, _, value = self.stack.pop() if pop else self.stack[-1]
        return value

    def _sx_peek_(self, pop: bool = True) -> T:
        value = self.pstack.pop() if pop else self.pstack[-1]
        return value

    def _sx_put_(self, value: T, pop: bool = True) -> None:
        context, segment, stream, pos, _ = self.stack.pop() if pop else self.stack[-1]
        with context.enter_segment(segment, stream, pos, os.SEEK_SET) as f:
            context.dump(to_type(self.child), f, value)

    def _sx_is_const_(self) -> bool:
        return False

    def __repr__(self) -> str:
        return f'{__name__}.TypeSource({super().__repr__()}, count={self.count!r})'