Three Dollar Quill

from typing import Generic, TypeVar

from wherever import algebraic

R = typing.TypeVar('R')
E = typing.TypeVar('E')

@algebraic
class Result(Generic[R, E]):

    Ok: R
    Err: E

"""Class decorator for defining abstract data types."""

import collections
import keyword
import typing


def _as_tuple(value) -> tuple:
    """Convert non-tuple values to a 1-tuple, return tuples unchanged."""
    if not isinstance(value, tuple):
        value = (value,)
    return value


def _name(cls, function) -> str:
    """Return the name of a function accessed through a descriptor."""
    return function.__get__(None, cls).__name__


def _set_new_functions(cls, *functions) -> typing.Optional[str]:
    """Attempt to set the attributes corresponding to the functions on cls.

    If any attributes are already defined, fail before setting, and return the
    already-defined name.
    """
    for function in functions:
        if _name(cls, function) in cls.__dict__:
            return _name(cls, function)
    for function in functions:
        setattr(cls, _name(cls, function), function)
    return None


class MatchFailure(BaseException):
    """An exception that signals a failure in ADT matching."""


def desugar(constructor: type, instance: tuple) -> tuple:
    """Return the inside of an ADT instance, given its constructor."""
    # I really do want to match exactly, I swear.
    if type(instance) is not constructor:
        raise MatchFailure
    return tuple.__getitem__(instance, slice(None))


def _unpack(instance: tuple) -> tuple:
    """Return the inside of any ADT instance.

    This function is not meant for general use.
    """
    return desugar(type(instance), instance)


def _algebraic_base(obj):
    return getattr(obj.__class__, '__algebraic_base__', None)


class AlgebraicConstructor:

    """Base class for ADT Constructor classes."""

    __slots__ = ()

    def __new__(cls, *args, **kwargs):
        """Explicitly forward to base class."""
        return super().__new__(cls, *args, **kwargs)


def __repr__(self):
    return self.__class__.__qualname__ + (
        f'({", ".join(repr(item) for item in _unpack(self))})')


def __eq__(self, other):
    if other.__class__ is self.__class__:
        return _unpack(self) == _unpack(other)
    if _algebraic_base(other) is _algebraic_base(self):
        return False
    return NotImplemented


def __ne__(self, other):
    if other.__class__ is self.__class__:
        return _unpack(self) != _unpack(other)
    if _algebraic_base(other) is _algebraic_base(self):
        return True
    return NotImplemented


def __lt__(self, other):
    if other.__class__ is self.__class__:
        return _unpack(self) < _unpack(other)
    if _algebraic_base(other) is _algebraic_base(self):
        order = _algebraic_base(self).__subclass_order__
        self_index = order.index(self.__class__)
        other_index = order.index(other.__class__)
        return self_index < other_index
    return NotImplemented


def __le__(self, other):
    if other.__class__ is self.__class__:
        return _unpack(self) <= _unpack(other)
    if _algebraic_base(other) is _algebraic_base(self):
        order = _algebraic_base(self).__subclass_order__
        self_index = order.index(self.__class__)
        other_index = order.index(other.__class__)
        return self_index <= other_index
    return NotImplemented


def __gt__(self, other):
    if other.__class__ is self.__class__:
        return _unpack(self) > _unpack(other)
    if _algebraic_base(other) is _algebraic_base(self):
        order = _algebraic_base(self).__subclass_order__
        self_index = order.index(self.__class__)
        other_index = order.index(other.__class__)
        return self_index > other_index
    return NotImplemented


def __ge__(self, other):
    if other.__class__ is self.__class__:
        return _unpack(self) >= _unpack(other)
    if _algebraic_base(other) is _algebraic_base(self):
        order = _algebraic_base(self).__subclass_order__
        self_index = order.index(self.__class__)
        other_index = order.index(other.__class__)
        return self_index >= other_index
    return NotImplemented


def __hash__(self):
    return hash(_unpack(self))


def _make_constructor(_cls, name, length, subclasses, subclass_order):
    class Constructor(_cls, AlgebraicConstructor, tuple):
        """Auto-generated subclass of an ADT."""
        __slots__ = ()

        __algebraic_base__ = _cls

        def __new__(cls, *args):
            if len(args) != length:
                raise ValueError
            return super().__new__(cls, args)

    Constructor.__name__ = name
    Constructor.__qualname__ = f'{_cls.__qualname__}.{name}'

    subclasses.add(Constructor)
    setattr(_cls, name, Constructor)
    subclass_order.append(Constructor)


def _process_class(_cls, _repr, eq, order):
    if order and not eq:
        raise ValueError('eq must be true if order is true')

    lengths = {}
    subclasses = set()
    subclass_order = []
    for cls in reversed(_cls.__mro__):
        for key, value in getattr(cls, '__annotations__', {}).items():
            lengths[key] = len(_as_tuple(value))

    for name, length in lengths.items():
        _make_constructor(_cls, name, length, subclasses, subclass_order)

    @classmethod
    def __init_subclass__(cls, **kwargs):
        if issubclass(cls, tuple(subclasses)):
            raise TypeError
        # Allow it to go through otherwise, because Generic.
        return super(_cls, cls).__init_subclass__(**kwargs)

    _cls.__init_subclass__ = __init_subclass__

    @staticmethod
    def __new__(cls, args):
        if cls not in subclasses:
            raise TypeError
        return super(_cls, cls).__new__(cls, args)

    if _set_new_functions(_cls, __new__):
        base__new__ = _cls.__new__

        @staticmethod
        def __new__(cls, args):
            if cls not in subclasses:
                raise TypeError
            return base__new__(cls, args)

        _cls.__new__ = __new__

    if _repr:
        _set_new_functions(_cls, __repr__)

    equality_methods_were_set = False

    if eq:
        equality_methods_were_set = not _set_new_functions(
            _cls, __eq__, __ne__)

    if equality_methods_were_set:
        _cls.__hash__ = __hash__

    if order:
        if not equality_methods_were_set:
            raise ValueError(
                "Can't add ordering methods if equality methods are provided.")
        collision = _set_new_functions(_cls, __lt__, __le__, __gt__, __ge__)
        if collision:
            raise TypeError(f'Cannot overwrite attribute {collision} '
                            f'in class {_cls.__name__}. Consider using '
                            'functools.total_ordering')

    _cls.__subclass_order__ = tuple(subclass_order)

    return _cls


def algebraic(_cls=None, *, repr=True, eq=True, order=False):
    """Decorate a class to be an algebraic data type."""

    def wrap(cls):
        """Return the processed class."""
        return _process_class(cls, repr, eq, order)

    if _cls is None:
        return wrap

    return wrap(_cls)


DISCARD = object()


class Matcher(tuple):
    """A matcher that binds a value to a name."""

    __slots__ = ()

    def __new__(cls, name: str):
        if name == '_':
            return DISCARD
        if not name.isidentifier():
            raise ValueError
        if keyword.iskeyword(name):
            raise ValueError
        return super().__new__(cls, (name,))

    @property
    def name(self):
        """Return the name of the matcher."""
        return self[0]

    def __matmul__(self, other):
        return AsMatcher(self, other)


class AsMatcher(tuple):
    """A matcher that contains further bindings."""

    __slots__ = ()

    def __new__(cls, matcher: Matcher, match):
        if matcher is DISCARD:
            return match
        return super().__new__(cls, (matcher, match))

    @property
    def matcher(self):
        """Return the left-hand-side of the as-match."""
        return self[0]

    @property
    def match(self):
        """Return the right-hand-side of the as-match."""
        return self[1]


def names(target):
    """Return every name bound by a target."""
    name_list = []
    names_seen = set()
    to_process = [target]
    while to_process:
        item = to_process.pop()
        if isinstance(item, Matcher):
            if item.name in names_seen:
                raise ValueError
            names_seen.add(item.name)
            name_list.append(item.name)
        elif isinstance(item, AsMatcher):
            to_process.append(item.match)
            to_process.append(item.matcher)
        elif isinstance(item, AlgebraicConstructor):
            to_process.extend(reversed(_unpack(item)))
        elif isinstance(item, tuple):
            to_process.extend(reversed(item))
    yield from name_list


def _match(target, value):
    match_dict = collections.OrderedDict()
    to_process = [(target, value)]
    while to_process:
        target, value = to_process.pop()
        if target is DISCARD:
            pass
        elif isinstance(target, Matcher):
            if target.name in match_dict:
                raise ValueError
            match_dict[target.name] = value
        elif isinstance(target, AsMatcher):
            to_process.append((target.match, value))
            to_process.append((target.matcher, value))
        elif isinstance(target, AlgebraicConstructor):
            to_process.extend(zip(reversed(_unpack(target)),
                                  reversed(desugar(type(target), value))))
        elif (isinstance(target, tuple) and
              target.__class__ is value.__class__ and
              len(target) == len(value)):
            to_process.extend(zip(reversed(target), reversed(value)))
        elif target != value:
            raise MatchFailure
    return match_dict


def get_values(dct, keys):
    """Unpack a dict, in order."""
    for key in keys:
        yield dct[key]


class ValueMatcher:
    """Given a value, attempt to match against a target."""

    def __init__(self, value):
        self.value = value
        self.matches = None

    def match(self, target):
        """Match against target, generating a set of bindings."""
        try:
            self.matches = _match(target, self.value)
        except MatchFailure:
            self.matches = None
        return self.matches is not None


if __name__ == '__main__':

    R = typing.TypeVar('R')
    E = typing.TypeVar('E')

    @algebraic
    class Result(typing.Generic[R, E]):
        """Experimental generic ADT."""

        Ok: R
        Err: E

    my_result: Result[int, str] = Result.Ok(10)