cjw296 / generative_typing Goto Github PK

from dataclasses import dataclass
from typing import TypedDict

bag_bad_init: Bag[int] = Bag[int](23.0)
bag_bad_var_type: Bag[int] = Bag[float](23.0)


@dataclass
class MyDataClass:
    foo: int
    bar: str


bag_dataclass: Bag[MyDataClass] = Bag[MyDataClass](MyDataClass(foo=1, bar='x'))
# Happy path examples
bag_dataclass_foo: Bag[int] = bag_dataclass.foo
foo: int = bag_dataclass.foo.get()
foo_: int = bag_dataclass_foo.get()
bag_dataclass_bar: Bag[str] = bag_dataclass.bar
bar: str = bag_dataclass.bar.get()
bar_: str = bag_dataclass_bar.get()
# unhappy
foo_bad: str = bag_dataclass.foo.get()
bar_bad: int = bag_dataclass.bar.get()


class MyTypedDict(TypedDict):
    foo: int
    bar: str


bag_typeddict: Bag[MyTypedDict] = Bag[MyTypedDict]({'foo': 1, 'bar': 'x'})
# Happy path examples
bag_typeddict_foo: Bag[int] = bag_typeddict['foo']
dfoo: int = bag_typeddict['foo'].get()
dfoo_: int = bag_typeddict_foo.get()
bag_typeddict_bar: Bag[str] = bag_typeddict['bar']
dbar: str = bag_typeddict['bar'].get()
dbar_: str = bag_typeddict_bar.get()
# unhappy
dfoo_bad: str = bag_typeddict['foo'].get()
dbar_bad: int = bag_typeddict['bar'].get()
bag_typeddict_bad: Bag[MyTypedDict] = Bag[MyTypedDict]({'foo': 'x', 'bar': 1})
dfoo_bad2 = bag_typeddict.foo  # no attribute!
bag_typeddict_foo_bad: Bag[str] = bag_typeddict['foo']


def bad_callable() -> int:
    return 1


Bag[MyDataClass](bad_callable)

bag_list: Bag[list[str]] = Bag[list[str]](['x', 'y', 'z'])
bag_list_1: Bag[str] = bag_list[1]
value: str = bag_list_1.get()
bag_list_bad_index: Bag[str] = bag_list['foo']  # index not int!
bag_list_bad_type: Bag[float] = bag_list['x']  # variable should be Bag[str]
bad_list_bad_type: int = bag_list['x'].get()  # variable should be str

food: dict[int, str] = {}
food['1'] = 2
# dict of strings

from mypy.errorcodes import ATTR_DEFINED, INDEX
from mypy.messages import format_type_inner
from mypy.nodes import MemberExpr, SymbolTableNode
from mypy.plugin import AttributeContext, MethodSigContext, Plugin
from mypy.types import FunctionLike, Instance, Type, TypeAliasType, TypedDictType


...

def attribute_type_from_generic_parent(ctx: AttributeContext) -> Type:
    info = ctx.type
    concrete_type = info.args[0]
    assert isinstance(info, Instance), type(info)
    attr_context = ctx.context
    assert isinstance(attr_context, MemberExpr), type(attr_context)
    attr_name = attr_context.name
    if isinstance(concrete_type, Instance):
        arg_info = concrete_type.type.names[attr_name]
        assert isinstance(arg_info, SymbolTableNode), type(arg_info)
        assert arg_info.type is not None
        attr_type = arg_info.type
    else:
        ctx.api.fail(
            f'{format_type_inner(info, verbosity=0, fullnames=set())} has no attribute {attr_name}',
            ctx.context,
            code=ATTR_DEFINED,
        )
        return ctx.default_attr_type
    return Instance(info.type, [attr_type])


def item_type_from_generic_parent(ctx: MethodSigContext) -> FunctionLike:
    info = ctx.type
    concrete_type = info.args[0]
    index = ctx.args[0][0].value
    if isinstance(concrete_type, Instance):
        concrete_info = concrete_type.type
        if concrete_info.fullname == 'builtins.list':
            import pdb; pdb.set_trace()
            ctx.default_signature.ret_type = Instance(info.type, [concrete_type.args[0]])
        else:
            assert False
    elif isinstance(concrete_type, TypeAliasType):
        assert concrete_type.alias is not None
        dict_type = concrete_type.alias.target
        assert isinstance(dict_type, TypedDictType)
        attr_type = dict_type.items[index]
        ctx.default_signature.ret_type = Instance(info.type, [attr_type])
    else:
        ctx.api.fail(
            f'{format_type_inner(info, verbosity=0, fullnames=set())} has no item {index}',
            ctx.context,
            code=INDEX,
        )
    return ctx.default_signature


...
    def get_attribute_hook(self, fullname: str) -> Optional[Callable[[AttributeContext], Type]]:
        if fullname.startswith('bag.Bag.'):
            attr = fullname.rsplit('.', 1)[-1]
            if not attr.startswith('_'):
                return attribute_type_from_generic_parent
        return None

    def get_method_signature_hook(self, fullname: str) -> Optional[Callable[[MethodSigContext], FunctionLike]]:
        if fullname == 'bag.Bag.__getitem__':
            return item_type_from_generic_parent
        return None