116 lines
4.6 KiB
Python
116 lines
4.6 KiB
Python
from typing import Union as U, Generic as G, Optional as O, TypeVar, Mapping, Sequence
|
|
from ..core.base import Type, Context, PathElement
|
|
from ..core.io import Stream, Pos, PosInfo
|
|
from ..core.util import format_value
|
|
from ..core.expr import Expr
|
|
from ..core import to_type
|
|
from .data import Nothing
|
|
|
|
T = TypeVar('T')
|
|
V = TypeVar('V')
|
|
|
|
class Switch(G[T, V], Type[T]):
|
|
def __init__(self, options: Mapping[V, Type[T]], selector: O[V] = None, default: O[V] = None, fallback: O[Type[T]] = None) -> None:
|
|
self.options = options
|
|
self.default_key = default
|
|
self.default_val = fallback
|
|
self.selector = selector
|
|
|
|
def get_value(self, context: Context, peek: bool = False) -> Type[T]:
|
|
default = context.peek(self.default_val) if peek else context.get(self.default_val)
|
|
options = context.peek(self.options) if peek else context.get(self.options)
|
|
if self.selector is not None:
|
|
selector = context.peek(self.selector) if peek else context.get(self.selector)
|
|
elif self.default_key is not None:
|
|
selector = context.peek(self.default_key) if peek else context.get(self.default_key)
|
|
else:
|
|
return to_type(default)
|
|
if selector not in options:
|
|
if default:
|
|
return to_type(default)
|
|
raise ValueError(f'selector {selector} not any in options {", ".join(str(o) for o in options)}')
|
|
return to_type(options[selector])
|
|
|
|
def parse(self, context: Context, stream: Stream) -> T:
|
|
child = self.get_value(context)
|
|
with context.enter(None, child):
|
|
return context.parse(child, stream)
|
|
|
|
def dump(self, context: Context, stream: Stream, value: T) -> None:
|
|
child = self.get_value(context)
|
|
with context.enter(None, child):
|
|
context.dump(child, stream, value)
|
|
|
|
def sizeof(self, context: Context, start: PosInfo, value: O[T]) -> O[PosInfo]:
|
|
child = self.get_value(context, peek=True)
|
|
with context.enter(None, child):
|
|
return context.sizeof(child, start, value)
|
|
|
|
def offsetof(self, context: Context, start: PosInfo, path: Sequence[PathElement], value: O[T]) -> O[PosInfo]:
|
|
child = self.get_value(context, peek=True)
|
|
with context.enter(None, child):
|
|
return context.offsetof(child, start, value)
|
|
|
|
def default(self, context: Context) -> T:
|
|
child = self.get_value(context, peek=True)
|
|
with context.enter(None, child):
|
|
return context.default(child)
|
|
|
|
def __str__(self) -> str:
|
|
return f'{format_value(self.options, str)}[{self.selector}]'
|
|
|
|
def __repr__(self) -> str:
|
|
return f'{__name__}.Switch({self.options!r}, selector={self.selector!r}, default={self.default_key!r}, fallback={self.default_val!r})'
|
|
|
|
class If(G[T,V], Type[U[T, V]]):
|
|
def __init__(self, cond: Expr, true: Type[T], false: Type[V] = Nothing()) -> None:
|
|
self.cond = cond
|
|
self.true = true
|
|
self.false = false
|
|
|
|
def parse(self, context: Context, stream: Stream) -> U[T, V]:
|
|
if context.get(self.cond):
|
|
child = to_type(self.true)
|
|
else:
|
|
child = to_type(self.false)
|
|
with context.enter(None, child):
|
|
return context.parse(child, stream)
|
|
|
|
def dump(self, context: Context, stream: Stream, value: U[T, V]) -> None:
|
|
if context.get(self.cond):
|
|
child = to_type(self.true)
|
|
else:
|
|
child = to_type(self.false)
|
|
with context.enter(None, child):
|
|
return context.dump(child, stream, value)
|
|
|
|
def sizeof(self, context: Context, start: PosInfo, value: O[U[T, V]]) -> O[PosInfo]:
|
|
if context.peek(self.cond):
|
|
child = to_type(self.true)
|
|
else:
|
|
child = to_type(self.false)
|
|
with context.enter(None, child):
|
|
return context.sizeof(child, start, value)
|
|
|
|
def offsetof(self, context: Context, start: PosInfo, path: Sequence[PathElement], value: O[U[T, V]]) -> O[PosInfo]:
|
|
if context.peek(self.cond):
|
|
child = to_type(self.true)
|
|
else:
|
|
child = to_type(self.false)
|
|
with context.enter(None, child):
|
|
return context.offsetof(child, start, path, value)
|
|
|
|
def default(self, context: Context) -> U[T, V]:
|
|
if context.peek(self.cond):
|
|
child = to_type(self.true)
|
|
else:
|
|
child = to_type(self.false)
|
|
with context.enter(None, child):
|
|
return context.default(child)
|
|
|
|
def __str__(self) -> str:
|
|
return f'({self.cond} ? {self.true} : {self.false})'
|
|
|
|
def __repr__(self) -> str:
|
|
return f'{__name__}.If({self.cond!r}, {self.true!r}, {self.false!r})'
|