Ibis Typehint Coercion supporting generic types

typehints.py

from typing import Sequence, Mapping, TypeVar

from ibis.common.grounds import Annotable
from ibis.common.validators import Coercible


T = TypeVar('T')
K = TypeVar('K')
V = TypeVar('V')


class List(Annotable, Sequence[T], Coercible):
    @classmethod
    def __coerce__(self, values):
        values = tuple(values)
        if values:
            head, *rest = values
            return ConsList(head, rest)
        else:
            return EmptyList()


class EmptyList(List[T]):
    def __getitem__(self, key):
        raise IndexError(key)

    def __len__(self):
        return 0


class ConsList(List[T]):
    head: T
    rest: List[T]

    def __getitem__(self, key):
        if key == 0:
            return self.head
        else:
            return self.rest[key - 1]

    def __len__(self):
        return len(self.rest) + 1


class Map(Annotable, Mapping[K, V], Coercible):
    @classmethod
    def __coerce__(self, pairs):
        pairs = dict(pairs)
        if pairs:
            head_key = next(iter(pairs))
            head_value = pairs.pop(head_key)
            rest = pairs
            return ConsMap((head_key, head_value), rest)
        else:
            return EmptyMap()


class EmptyMap(Map[K, V]):
    def __getitem__(self, key):
        raise KeyError(key)

    def __iter__(self):
        return iter(())

    def __len__(self):
        return 0


class ConsMap(Map[K, V]):
    head: Tuple[K, V]
    rest: Map[K, V]

    def __getitem__(self, key):
        if key == self.head[0]:
            return self.head[1]
        else:
            return self.rest[key]

    def __iter__(self):
        yield self.head[0]
        yield from self.rest

    def __len__(self):
        return len(self.rest) + 1


class Integer(int, Coercible):
    @classmethod
    def __coerce__(cls, value):
        return Integer(value)


class Float(float, Coercible):
    @classmethod
    def __coerce__(cls, value):
        return Float(value)


class MyExpr(Annotable):
    a: Integer
    b: List[Float]
    c: Map[str, Integer]


expr = MyExpr("1", ["2.0", 3, True], c={"a": "1", "b": 2, "c": 3.0})
assert expr.a == 1
assert expr.b == ConsList(2.0, ConsList(3.0, ConsList(1.0, EmptyList())))
assert expr.c == ConsMap(("a", 1), ConsMap(("b", 2), ConsMap(("c", 3), EmptyMap())))