hk.py

from abc import ABCMeta
from abc import abstractmethod
from dataclasses import dataclass
from typing import Any
from typing import Callable
from typing import Generic
from typing import Optional
from typing import TypeVar
from typing import cast

#####

A = TypeVar('A')
B = TypeVar('B')
C = TypeVar('C')
R = TypeVar('R')
X = TypeVar('X')

#####

class TyEq(Generic[A, B], metaclass=ABCMeta):
  @abstractmethod
  def cast(self, a: A) -> B: ...

  @abstractmethod
  def sym(self) -> TyEq[B, A]: ...

  @abstractmethod
  def compose(self, other: TyEq[B, C]) -> TyEq[A, C]: ...

@dataclass(frozen=True)
class TyEq_Refl(Generic[A], TyEq[A, A]):
  def cast(self, a: A) -> A:
    return a

  def sym(self) -> TyEq[A, A]:
    return self

  def compose(self, other: TyEq[A, C]) -> TyEq[A, C]:
    return other

TyEq_Refl[int]()

#####

N = TypeVar('N', bound='Nat')
M = TypeVar('M', bound='Nat')
XN = TypeVar('XN', bound='Nat')

class Nat(Generic[N], metaclass=ABCMeta):
  @abstractmethod
  def accept(self, visitor: Nat_Visitor[R, N]) -> R: ...

  @abstractmethod
  def eq(self, other: M) -> Optional[TyEq[N, M]]: ...

  @staticmethod
  def eq_succ(teq: TyEq[N, M]) -> TyEq[Nat_S[N], Nat_S[M]]:
    # TODO
    return cast(Any, TyEq_Refl())

@dataclass(frozen=True)
class Nat_Z(Nat[Nat_Z]):
  def accept(self, visitor: Nat_Visitor[R, Nat_Z]) -> R:
    return visitor.visit_z(TyEq_Refl(), self)

  def eq(self, other: M) -> Optional[TyEq[Nat_Z, M]]:
    class _Match(Nat_Visitor[Optional[TyEq[Nat_Z, M]], M]):
      def visit_z(self, teq: TyEq[M, Nat_Z], z: Nat_Z) -> Optional[TyEq[Nat_Z, M]]:
        return teq.sym()

      def visit_s(self, teq: TyEq[M, Nat_S[XN]], s: Nat_S[XN]) -> Optional[TyEq[Nat_Z, M]]:
        return None

    return other.accept(_Match())

@dataclass(frozen=True)
class Nat_S(Generic[N], Nat[Nat_S[N]]):

  prev: N

  def accept(self, visitor: Nat_Visitor[R, Nat_S[N]]) -> R:
    return visitor.visit_s(TyEq_Refl(), self)

  def eq(self, other: M) -> Optional[TyEq[Nat_S[N], M]]:
    class _Match(Nat_Visitor[Optional[TyEq[Nat_S[N], M]], M]):
      def visit_z(self2, teq: TyEq[M, Nat_Z], z: Nat_Z) -> Optional[TyEq[Nat_S[N], M]]:
        return None

      def visit_s(self2, teq: TyEq[M, Nat_S[XN]], s: Nat_S[XN]) -> Optional[TyEq[Nat_S[N], M]]:
        eqNX: Optional[TyEq[N, XN]] = self.prev.eq(s.prev)
        if eqNX is None:
          return None
        else:
          return Nat.eq_succ(eqNX).compose(teq.sym())

    return other.accept(_Match())

class Nat_Visitor(Generic[R, N], metaclass=ABCMeta):
  @abstractmethod
  def visit_z(self, teq: TyEq[N, Nat_Z], z: Nat_Z) -> R: ...

  @abstractmethod
  def visit_s(self, teq: TyEq[N, Nat_S[M]], s: Nat_S[M]) -> R: ...

two: Nat[Nat_S[Nat_S[Nat_Z]]] = Nat_S(Nat_S(Nat_Z()))

#####

W = TypeVar('W')

class Hk(Generic[W, A], metaclass=ABCMeta):
  pass

#####

class Functor(Generic[W], metaclass=ABCMeta):
  @abstractmethod
  def map(self, fa: Hk[W, A], f: Callable[[A], B]) -> Hk[W, B]: ...

class Applicative(Generic[W], metaclass=ABCMeta):
  @abstractmethod
  def pure(self, a: A) -> Hk[W, A]: ...

  @abstractmethod
  def liftA2(self, f: Callable[[A, B], C], fa: Hk[W, A], fb: Hk[W, B]) -> Hk[W, C]: ...

#####

class Maybe(Generic[A], Hk[Maybe_Mu, A], metaclass=ABCMeta):
  @staticmethod
  def narrow(hk: Hk[Maybe_Mu, A]) -> Maybe[A]:
    return cast(Maybe[A], hk)

class Maybe_Mu(metaclass=ABCMeta):
  pass

@dataclass(frozen=True)
class Nothing(Generic[A], Maybe[A]):
  pass

@dataclass(frozen=True)
class Just(Generic[A], Maybe[A]):
  value: A

@dataclass
class Maybe_Functor(Functor[Maybe_Mu]):
  def map(self, fa: Hk[Maybe_Mu, A], f: Callable[[A], B]) -> Hk[Maybe_Mu, B]:
    ma = Maybe.narrow(fa)
    if isinstance(ma, Just):
      return Just(f(ma.value))
    else:
      return Nothing()

@dataclass
class Maybe_Applicative(Applicative[Maybe_Mu]):
  def pure(self, a: A) -> Hk[Maybe_Mu, A]:
    return Just(a)

  def liftA2(self, f: Callable[[A, B], C], fa: Hk[Maybe_Mu, A], fb: Hk[Maybe_Mu, B]) -> Hk[Maybe_Mu, C]:
    ma = Maybe.narrow(fa)
    mb = Maybe.narrow(fb)
    if isinstance(ma, Just) and isinstance(mb, Just):
      return Just(f(ma.value, mb.value))
    else:
      return Nothing()

#####