sx/sx/core/base.py

import os
from types import SimpleNamespace

from contextlib import contextmanager
from typing import Any, Generic, Generator, Iterable, List, Mapping, Dict, Sequence, Tuple, TypeVar, Union as U, Optional as O, Generic as G, cast

from .util import seeking
from .io import Segment, Stream, Pos


class Params:
    __slots__ = ('segments', 'default_segment', 'user')

    def __init__(self, segments: Sequence[Segment] = None):
        default = segments[0] if segments else Segment('default')
        self.segments = {s.name: s for s in (segments or [default, Segment('refs', [default])])}
        self.default_segment = default
        self.user = SimpleNamespace()

    def reset(self):
        for s in self.segments.values():
            s.reset()


PathElement = U[str, int]
PathEntry = Tuple[PathElement, 'Type']

def format_path(path: Iterable[PathElement]) -> str:
    s = ''
    first = True
    for p in path:
        sep = '.'
        if isinstance(p, int):
            p = '[' + str(p) + ']'
            sep = ''
        if sep and not first:
            s += sep
        s += p
        first = False
    return s


T = TypeVar('T')
PT = TypeVar('PT')


class PossibleDynamic(Generic[T]):
    pass

class Context:
    __slots__ = ('root', 'value', 'params', 'path', 'segment_path')

    def __init__(self, root: 'Type', value: O[Any] = None, params: O[Params] = None) -> None:
        self.root = root
        self.value = value
        self.params = params or Params()
        self.path: List[PathEntry] = []
        self.segment_path: List[Segment] = []

    def copy(self) -> 'Context':
        c = self.__class__(root=self.root, value=self.value, params=self.params)
        c.path = self.path.copy()
        c.segment_path = self.segment_path.copy()
        return c


    @property
    def segment(self) -> Segment:
        return self.segment_path[-1] if self.segment_path else self.params.default_segment

    @contextmanager
    def enter(self, entry: PathElement, type: 'Type') -> Generator:
        self.path.append((entry, type))
        try:
            yield
        except EOFError as e:
            raise EOF(self, e) from e
        self.path.pop()

    @contextmanager
    def enter_segment(self, segment: Segment, stream: O[Stream] = None, pos: O[Pos] = None, reference = os.SEEK_SET) -> Generator[O[Stream], None, None]:
        if stream:
            if pos is None:
                if segment.offset is None:
                    segment.offset = self.segment_offset(segment)
                    segment.pos = segment.offset
                pos = segment.pos
            if pos is None:
                raise Error(self, ValueError('could not enter segment {}: could not calculate offset'.format(segment)))
            with seeking(stream.root, pos, reference) as s, stream.wrapped(s) as f:
                self.segment_path.append(segment)
                yield f
                self.segment_path.pop()
                segment.pos = f.tell()
        else:
            self.segment_path.append(segment)
            yield stream
            self.segment_path.pop()

    def segment_offset(self, segment: Segment) -> O[Pos]:
        size: Pos = 0
        for s in segment.dependents:
            sz = self.segment_size(s)
            if sz is None:
                return None
            off = self.segment_offset(s)
            if off is None:
                return None
            size += off + sz
        return size

    def segment_size(self, segment: Segment) -> O[Pos]:
        sizes = self.sizeof(self.root, self.value)
        return sizes.get(segment, None)

    def format_path(self) -> str:
        return format_path(name for name, _ in self.path)

    def to_size(self, value: Any) -> Dict[Segment, Pos]:
        if not isinstance(value, dict):
            stream = self.segment_path[-1] if self.segment_path else self.params.default_segment
            value = {stream: value}
        return value


    def get(self, value: U[T, PossibleDynamic[T]]) -> T:
        from .expr import Expr, get
        if isinstance(value, Expr):
            value = get(value)
        return cast(T, value)

    def peek(self, value: U[T, PossibleDynamic[T]]) -> O[T]:
        from .expr import Expr, peek
        if isinstance(value, Expr):
            value = peek(value)
        return cast(T, value)

    def put(self, value: U[T, PossibleDynamic[T]], new: T) -> None:
        from .expr import Expr, put
        if isinstance(value, Expr):
            put(value, new)


    def parse(self, type: 'Type[PT]', stream: Stream) -> PT:
        return type.parse(self, stream)

    def dump(self, type: 'Type[PT]', stream: Stream, value: PT) -> None:
        return type.dump(self, stream, value)

    def sizeof(self, type: 'Type[PT]', value: O[PT] = None) -> Dict[Segment, Pos]:
        return self.to_size(type.sizeof(self, value))

    def offsetof(self, type: 'Type[PT]', path: Sequence[PathElement], value: O[PT] = None) -> Dict[Segment, Pos]:
        return self.to_size(type.offsetof(self, path, value))

    def default(self, type: 'Type[PT]') -> PT:
        return type.default(self)


class Type(G[PT]):
    __slots__ = ()

    def parse(self, context: Context, stream: Stream) -> PT:
        raise NotImplementedError

    def dump(self, context: Context, stream: Stream, value: PT) -> None:
        raise NotImplementedError

    def sizeof(self, context: Context, value: O[PT]) -> U[Mapping[str, int], O[int]]:
        return None

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

    def default(self, context: Context) -> PT:
        raise NotImplementedError


class Error(Exception):
    __slots__ = ('context',)

    def __init__(self, context: Context, exception: Exception) -> None:
        path = context.format_path()
        if path:
            path = '[' + path + '] '
        if not isinstance(exception, Exception):
            exception = ValueError(exception)

        super().__init__('{}{}: {}'.format(
            path, exception.__class__.__name__, str(exception),
        ))
        self.exception = exception
        self.context = context.copy()

class EOF(Error):
    pass