control: move Switch from transforms and add If type

This commit is contained in:
Shiz 2021-06-26 05:30:11 +02:00
parent fd7476a434
commit 02dcfe0d92
3 changed files with 97 additions and 42 deletions

View File

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

93
sx/types/control.py Normal file
View File

@ -0,0 +1,93 @@
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
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 and default:
return to_type(default)
return to_type(options[selector])
def parse(self, context: Context, stream: Stream) -> T:
return context.parse(self.get_value(context), stream)
def dump(self, context: Context, stream: Stream, value: T) -> None:
context.dump(self.get_value(context), stream, value)
def sizeof(self, context: Context, value: O[T]) -> O[Pos]:
return context.sizeof(self.get_value(context, peek=True), value)
def offsetof(self, context: Context, path: Sequence[PathElement], value: O[T]) -> O[Pos]:
return context.offsetof(self.get_value(context, peek=True), path, value)
def default(self, context: Context) -> T:
return context.default(self.get_value(context, peek=True))
def __str__(self) -> str:
return f'{format_value(self.options)}[{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):
return context.parse(to_type(self.true), stream)
else:
return context.parse(to_type(self.false), stream)
def dump(self, context: Context, stream: Stream, value: U[T, V]) -> None:
if context.get(self.cond):
context.dump(to_type(self.true), stream, value)
else:
context.dump(to_type(self.false), stream, value)
def sizeof(self, context: Context, value: O[U[T, V]]) -> O[Pos]:
if context.peek(self.cond):
return context.sizeof(to_type(self.true), value)
else:
return context.sizeof(to_type(self.false), value)
def offsetof(self, context: Context, path: Sequence[PathElement], value: O[U[T, V]]) -> O[Pos]:
if context.peek(self.cond):
return context.offsetof(to_type(self.true), path, value)
else:
return context.offsetof(to_type(self.false), path, value)
def default(self, context: Context) -> U[T, V]:
if context.peek(self.cond):
return context.default(to_type(self.true))
else:
return context.default(to_type(self.false))
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})'

View File

@ -209,43 +209,3 @@ class Enum(G[T, E], Transform[T, E]):
str=f'{enum.__name__}({child})',
repr=f'{__name__}.Enum({enum!r}, {child!r})',
)
class Switch(G[T, V], Type[T]):
def __init__(self, options: Mapped[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)
return to_type(options[selector])
def parse(self, context: Context, stream: Stream) -> T:
return context.parse(self.get_value(context), stream)
def dump(self, context: Context, stream: Stream, value: T) -> None:
context.dump(self.get_value(context), stream, value)
def sizeof(self, context: Context, value: O[T]) -> O[Pos]:
return context.sizeof(self.get_value(context, peek=True), value)
def offsetof(self, context: Context, path: Sequence[PathElement], value: O[T]) -> O[Pos]:
return context.offsetof(self.get_value(context, peek=True), path, value)
def default(self, context: Context) -> T:
return context.default(self.get_value(context, peek=True))
def __str__(self) -> str:
return f'{format_value(self.options)}[{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})'