generics.py

import sys
from functools import wraps
from inspect import isfunction, signature, _empty
from types import FunctionType
from typing import TypeVar, _GenericAlias, ForwardRef

try:
    from pytypes import is_of_type
    TYPE_CLS = (type, _GenericAlias)
except ImportError:
    is_of_type = isinstance
    TYPE_CLS = (type,)


def get_type_hints(x):
    ann = getattr(x, "__annotations__", {})

    if isinstance(x, type) or isfunction(x):
        # Class
        base_globals = sys.modules[x.__module__].__dict__
        for k, v in ann.items():
            if isinstance(v, str):
                try:
                    ann[k] = eval(v, base_globals)
                except NameError:
                    pass

    if ann:
        setattr(x, "__annotations__", ann)

    return ann


def _stringify_type(t):
    if isinstance(t, str):
        return t
    if isinstance(t, type):
        return t.__qualname__
    raise Exception(str(t))


def makeprop(name, typ, initial):
    def setter(_, val: typ):
        if isinstance(typ, type):
            assert is_of_type(val, typ), \
                f"Attempted to set property {name} of type {typ.__qualname__}" \
                f" to value of type {type(val).__qualname__}"

        props = getattr(_, "__generic_properties", {})
        props[name] = val
        setattr(_, "__generic_properties", props)

    def getter(_) -> typ:
        props = getattr(_, "__generic_properties", {})
        val = props.get(name, initial)
        setattr(_, "__generic_properties", props)

        assert val is not GenericWrapper.NO_VALUE, \
            f"Attempted to get property {name} but was never set"

        if isinstance(typ, type):
            assert isinstance(val, typ), \
                f"Property {name} of type {typ.__qualname__} turned into a value of type {type(val).__qualname__}"
        return val

    prop = property(getter, setter)
    return prop


def _remap(r_type, mapping):
    if isinstance(r_type, str):
        return mapping.get(r_type, r_type)

    if isinstance(r_type, _GenericAlias):
        r_type.__args__ = (*(
            _remap(x.__forward_arg__ if isinstance(x, ForwardRef) else x, mapping)
            for x in r_type.__args__
        ),)
        return r_type

    if not (hasattr(r_type, "__dict__") and "___ctor" in r_type.__dict__):
        return r_type

    ctor = r_type.__dict__['___ctor']
    args = ctor.args
    map = {**r_type.__dict__['___mapping']}

    for g, t in map.items():
        if t in mapping:
            map[g] = mapping[t]
        else:
            map[g] = _remap(t, mapping)

    items = (map.get(x, x) for x in args)

    return ctor[(*items,)]


def _wrap_func(func, mapping):
    func.__hints_done = False

    if func.__annotations__:  # We can do type checking
        @wraps(func)
        def wrapper(*args, **kwargs):
            if not func.__hints_done:
                na = {}
                for k, v in get_type_hints(func).items():
                    nv = _remap(v, mapping)
                    if isinstance(nv, _GenericAlias):
                        nv.__qualname__ = nv
                    na[k] = nv

                func.__annotations__ = na
                func.__hints_done = True

            sig = signature(func)  # Prepare sig

            # ====

            bound = sig.bind(*args, **kwargs)  # Bind arguments passed

            for k, v in sig.parameters.items():
                if v.kind == v.VAR_POSITIONAL:
                    # Tuple check
                    if isinstance(v.annotation, TYPE_CLS) and v.annotation is not _empty:
                        assert all(is_of_type(x, v.annotation) for x in bound.arguments.get(k, [])), \
                            (f"Parameter '{k}' was of type '{type(bound.arguments.get(k, [])).__qualname__}', "
                             f"expected '{v.annotation.__qualname__}'")

                elif v.kind == v.VAR_KEYWORD:
                    # dict check
                    print("WARNING: Unable to check var_keyword arguments (i.e. **kwargs)")
                else:
                    if isinstance(v.annotation, type) and v.annotation is not _empty:
                        assert is_of_type(bound.arguments[k], v.annotation), \
                            (f"Parameter '{k}' was of type '{type(bound.arguments[k]).__qualname__}', "
                             f"expected '{v.annotation.__qualname__}'")

            retval = func(*args, **kwargs)

            if isinstance(sig.return_annotation, TYPE_CLS) and sig.return_annotation is not _empty:
                assert is_of_type(retval, sig.return_annotation), \
                    f"Return type was of type '{type(retval).__qualname__}', expected " \
                    f"'{sig.return_annotation.__qualname__}'"

            return retval
        return wrapper
    return func


class GenericWrapper:
    CACHE = {}

    class NO_VALUE:
        pass

    def __init__(self, cls, args):
        assert isinstance(cls, type)
        self.cls = cls
        self.CACHE[cls.__qualname__] = self.CACHE.get(cls.__qualname__, {})
        self.bases = [*self.cls.__bases__]
        self.props = {**get_type_hints(self.cls)}
        self.attrs = {**self.cls.__dict__}
        self.args = args

    def copy(self):
        new = GenericWrapper(self.cls, (*self.args,))
        new.bases = [*self.bases]
        new.props = {**self.props}
        new.attrs = {**self.attrs}
        return new

    def __getitem__(self, item):
        # Everything below here should happen in a different class or make modifications to another object,
        # to prevent it leaking over into other instances
        assert isinstance(item, (tuple, type, str, TypeVar))
        types = item if isinstance(item, tuple) else (item,)

        if types in self.CACHE[self.cls.__qualname__]:
            return self.CACHE[self.cls.__qualname__][types]

        obj = self.copy()

        mapping = {k: v for k, v in zip(self.args, types)}
        mcs = getattr(self.cls, "__metaclass__", type)

        obj.remap_props(mapping)
        obj.remap_bases(mapping)
        obj.remap_attrs(mapping)

        cls = mcs(
            f"{obj.cls.__name__.split('[')[0]}[{', '.join(_stringify_type(mapping.get(v, v)) for v in obj.args)}]",
            (*obj.bases,),
            {**obj.attrs,
             "__annotations__": {**obj.props},

             "___mapping": mapping,
             "_": lambda s, p: mapping[p],
             "___ctor": obj}
        )
        self.CACHE[self.cls.__qualname__][types] = cls
        return cls

    def remap_props(self, mapping):
        for k, v in self.props.items():
            # k: prop name
            # v: type
            # mapping: map[old_type => new_type]
            t = _remap(v, mapping)
            if isinstance(t, str):
                continue

            self.props[k] = t

            initial_value = self.attrs.get(k, self.NO_VALUE)
            if isinstance(initial_value, property) and hasattr(initial_value.fget, "_generic_generated"):
                initial_value = initial_value.fget._generic_generated

            prop = makeprop(k, t, initial_value)

            self.attrs[k] = prop

    def remap_bases(self, mapping):
        self.bases = [
            _remap(v, mapping)
            for v in self.bases
        ]

    def remap_attrs(self, mapping):
        for k, v in self.attrs.items():
            # only methods
            if isfunction(v):
                v = getattr(v, "__wrapped__", v)
                nf = FunctionType(v.__code__, v.__globals__, v.__qualname__, v.__defaults__, v.__closure__)
                nf.__annotations__ = v.__annotations__
                self.attrs[k] = _wrap_func(nf, mapping)

    def __repr__(self):
        return f"<untyped generic class '{self.cls.__module__}.{self.cls.__qualname__}'>"


def generic(*args: str):
    def generic_inner(cls):
        return GenericWrapper(cls, args)
    return generic_inner