python / mypy

Optional static typing for Python
https://www.mypy-lang.org/
Other
18.21k stars 2.78k forks source link

How to make a class Unpackable from plugin hook. #10208

Open giannitedesco opened 3 years ago

giannitedesco commented 3 years ago

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?

ljluestc commented 3 months ago
# 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