mypy icon indicating copy to clipboard operation
mypy copied to clipboard

How to make a class Unpackable from plugin hook.

Open giannitedesco opened this issue 5 years ago • 1 comments

Hi guys,

I am trying to add mypy support for dataclassy, which is a replacement for standard dataclasses.

You can see the code here: https://github.com/biqqles/dataclassy/blob/15cef524b55df2c65b39d1619c7b8f1a3d28a813/dataclassy/mypy.py

One of the problems I've run in to is in supporting the feature of dataclassy that allows its dataclasses to be unpacked like tuples.

Firstly, I couldn't figure out how to represent tuple style unpacking in the mypy type system. If I create a class with def __iter__(self) -> Iterable[int], then it works, but I can't seem to do anything like Iterable[int, int, str, bool]. When I reveal_type for a NamedTuple.__iter__ I just see a generic iterator typing.Iterator[_T_co]. I see that reveal_type((1, 'foo').__iter__()) is typing.Iterator[builtins.object*]

So I tried to figure out if there was special logic in semanal_namedtuple.py but I couldn't see anything there other than adding Iterable as a base in the _make method(?) which doesn't seem like that I need. And, besides, I figure wherever the logic is it must also apply to regular tuples too.

What am I missing?

giannitedesco avatar Mar 13 '21 04:03 giannitedesco

# mypy_dataclassy_plugin.py

from mypy.plugin import Plugin, ClassDefContext
from mypy.types import TupleType, TypeOfAny
from mypy.nodes import TypeInfo, ARG_POS
from typing import Optional


class CustomDataClassPlugin(Plugin):
    def get_base_class_hook(self, fullname: str):
        if fullname == 'your_module.CustomDataClass':
            return add_tuple_iter
        return None


def add_tuple_iter(ctx: ClassDefContext) -> None:
    info = ctx.cls.info
    fallback = ctx.api.named_type('__builtins__.tuple')
    item_types = [ctx.api.named_type('__builtins__.int'),
                  ctx.api.named_type('__builtins__.int'),
                  ctx.api.named_type('__builtins__.str'),
                  ctx.api.named_type('__builtins__.bool')]
    iter_type = ctx.api.named_type('__builtins__.tuple')
    ctx.cls.info['__iter__'] = info['__iter__'] = make_callable_type(ctx, iter_type)


def make_callable_type(ctx: ClassDefContext, return_type) -> None:
    args = [ctx.api.named_type('__builtins__.self')]
    arg_kinds = [ARG_POS]
    return ctx.api.named_type('__builtins__.iter')
    fallback = ctx.api.named_type('__builtins__.tuple')
    return ctx.api.named_type('__builtins__.callable', [args, return_type, fallback])


def plugin(version: str):
    return CustomDataClassPlugin

ljluestc avatar Jun 17 '24 01:06 ljluestc