switch.py

__all__ = ['switch']


from contextlib import contextmanager
from typing import (
    cast,
    Callable,
    Generator,
    Generic,
    TypeVar,
    Union,
)


T = TypeVar('T')
GuardT = Callable[[T], bool]


class Case(Generic[T]):
    def __init__(self, value: T):
        self.value = value
        self.finished = False

    def __call__(self, cond: Union[T, GuardT]) -> bool:
        if (self.finished
                or (hasattr(cond, '__call__')
                    and not cast(GuardT, cond)(self.value))
                or cond != self.value):
            return False

        self.finished = True
        return True

    def fallthrough(self) -> None:
        self.finished = False

    def default(self) -> bool:
        return not self.finished


@contextmanager
def switch(value: T) -> Generator[Case[T], None, None]:
    yield Case(value)