Enhance `stubgen` to include types for `@dataclass` private methods

[Cireo]

Feature

Enhance stubgen to include types for @dataclass private methods.

Pitch

Was using interface files to help communicate between teams and noticed this (perhaps a low priority bug, but feels like a feature).

For example, currently

from dataclasses import dataclass

@dataclass(order=True)
class Dummy:
    foo: str
    bar: int

generates the following when passed to stubgen

from typing import Any

class Dummy:
    foo: str
    bar: int
    def __init__(self, foo: Any, bar: Any) -> None: ...
    def __lt__(self, other: Any) -> Any: ...
    def __gt__(self, other: Any) -> Any: ...
    def __le__(self, other: Any) -> Any: ...
    def __ge__(self, other: Any) -> Any: ...

I would hope we could generate

class Dummy:
    foo: str
    bar: int
    def __init__(self, foo: str, bar: bar) -> None: ...
    def __lt__(self, other: Dummy) -> bool: ...
    def __gt__(self, other: Dummy) -> bool: ...
    def __le__(self, other: Dummy) -> bool: ...
    def __ge__(self, other: Dummy) -> bool: ...

I believe there is already logic in the dataclass plugin to implement some of these checks, but if they are they are not propagated to stubgen.

Metadata

OSX
Python 3.9.0
mypy 0.790

[teskje]

At the moment, stubgen doesn't seem to handle dataclasses very well in general. I've looked into this a bit.

With the current master branch, if I generate stubs for this dataclass:

from dataclasses import dataclass
from typing import Optional

@dataclass
class Data:
    integer: int
    string: str
    optional: Optional[int] = None

It outputs:

from typing import Optional                                                                                                         
class Data:
    integer: int
    string: str
    optional: Optional[int]
    def __init__(self, integer, string, optional) -> None: ...

The issues are:

• Methods generated by the @dataclass decorator are missing type annotations
• The optional argument is not marked as a default argument (i.e. an = ... is missing)

At least the first issue (maybe the second too) is caused by the fact that __init__ is not present in the original code but generated by mypy.plugins.dataclasses. stubgen uses the original/unanalyzed code for adding type annotations to the functions it outputs (see https://github.com/python/mypy/blob/master/mypy/stubgen.py#L625-L626) but this doesn't exist in the case of these generated methods, so we end up without type annotations at all.

The generated methods still provide "analyzed" type information and I tried naively patching stubgen to use that, which resulted in:

from typing import Optional

class Data:
    integer: int
    string: str
    optional: Optional[int]
    def __init__(self, integer: builtins.int, string: builtins.str, optional: Union[builtins.int, None]) -> None: ...

Now there are type annotations and they are also correct, but the right imports (builtins and Union) are missing. I'm not sure if it is viable to try to add them. For this special case it seems easy enough, but what if we were dealing with custom classes defined in some other module?

Maybe a more feasible approach would be to simply leave the @dataclass decorator in the generated type stubs. If stubgen would output this:

from dataclasses import dataclass
from typing import Optional

@dataclass
class Data:
    integer: int
    string: str
    optional: Optional[int] = ...

... then everything should work nicely, no?

To get there, the following needs to change:

• stubgen needs to retain the @dataclass decorator. This seems easy enough to me.
• stubgen needs to stop emitting the generated dataclass methods (like __init__, __lt__, ...). Maybe there is a way to disable the mypy.plugin.dataclasses plugin when we analyze the code for stubgen?
• The = ... needs to be added for fields that have default values. This would require extra logic to differentiate between stuff like optional: int = field(default=1) and required: int = field(init=True).

[JelleZijlstra]

Thanks for the analysis! I also think leaving @dataclass in the stub is the right approach. Type checkers are expected to support @dataclass decorators in stubs.

[Cireo]

That approach works for me too, seems like a good separation of responsibilities.

[MinsungKim-BlockCrafters]

I tested following code with Mypy 0.961. The issues @teskje said still persists.

# test.py

from dataclasses import dataclass
from typing import Optional

@dataclass
class Data:
    required: int
    optional: Optional[int] = None

# test.pyi generated via stubgen

from typing import Optional

class Data:
    required: int
    optional: Optional[int]
    def __init__(self, required, optional) -> None: ...

By the way, why there is no --version parameter on stubgen?