batch list

BatchList.py

T = TypeVar('T')


class BatchList(List[T]):
    __slots__ = ()

    def __getattribute__(self, name):
        try:
            list_attr = list.__getattribute__(self, name)
            return list_attr
        except AttributeError:
            pass

        return BatchList(
            getattr(i, name) for i in self
        )

    def __call__(self, *a, **k):
        return BatchList(
            f(*a, **k) for f in self
        )

batchlist_mypyplugin.py

# -*- coding: utf-8 -*-

# -- stdlib --
from typing import Callable, Optional, cast

# -- third party --
from mypy.nodes import MemberExpr
from mypy.plugin import AttributeContext, MethodSigContext, Plugin
from mypy.types import AnyType, CallableType, Instance, Type, TypeOfAny

# -- own --

# -- code --


class BatchListPlugin(Plugin):

    def get_method_signature_hook(self, fullname: str
                                  ) -> Optional[Callable[[MethodSigContext], CallableType]]:
        print('method sig', fullname)
        if fullname == 'utils.misc.BatchList.__call__':
            return batchlist_method_signature_hook

        return None

    def get_attribute_hook(self, fullname: str
                           ) -> Optional[Callable[[AttributeContext], Type]]:
        if fullname.startswith('utils.misc.BatchList'):
            return batchlist_attribute_hook

        return None


def batchlist_attribute_hook(ctx: AttributeContext) -> Type:
    instance = cast(Instance, ctx.type)
    typeinfo = instance.type

    expr = cast(MemberExpr, ctx.context)
    field = expr.name

    if field in typeinfo.names:
        t = typeinfo.names[field].type
        assert t
        return t

    names = cast(Instance, instance.args[0]).type.names

    if field not in names:
        ctx.api.fail(
            'BatchList item {} has no attribute "{}"'.format(
                instance.args[0], field,
            ), expr
        )
        return Instance(typeinfo, [AnyType(TypeOfAny.from_error)])

    t = names[field].type
    assert t
    return Instance(typeinfo, [t])


def batchlist_method_signature_hook(ctx: MethodSigContext) -> CallableType:
    instance = cast(Instance, ctx.type)
    typeinfo = instance.type

    expr = ctx.context

    c = instance.args[0]
    if isinstance(c, CallableType):
        return c.copy_modified(
            ret_type=Instance(typeinfo, [c.ret_type])
        )
    else:
        ctx.api.fail("BatchList item {} is not callable".format(instance.args[0]), expr)
        return ctx.default_signature


def plugin(version: str) -> type:
    return BatchListPlugin