Skip to content

Commit

Permalink
stubgen: Fix generated dataclass __init__ signature (#16906)
Browse files Browse the repository at this point in the history
Fixes #16811

stubgen was swallowing default values for `__init__` methods generated
by the dataclass plugin making their signature incorrect. This is
because the plugin does not include the argument's initializer in the
generated signature. I changed it to include a dummy ellipsis so that
stubgen can generate correct code.
I also fixed arguments added by the dataclass plugin with the invalid
names `*` and `**` to have the valid and unique names `*generated_args`
and `**generated_kwargs` (with extra underscores to make them unique if
necessary). This removes the need for the hack to special case them in
stubgen and is less confusing for someone looking at them in a stub
file.
  • Loading branch information
hamdanal authored Feb 16, 2024
1 parent 2e5174c commit bfbac5e
Show file tree
Hide file tree
Showing 4 changed files with 50 additions and 31 deletions.
17 changes: 12 additions & 5 deletions mypy/plugins/dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
Context,
DataclassTransformSpec,
Decorator,
EllipsisExpr,
Expression,
FuncDef,
FuncItem,
Expand Down Expand Up @@ -149,13 +150,13 @@ def to_argument(
return Argument(
variable=self.to_var(current_info),
type_annotation=self.expand_type(current_info),
initializer=None,
initializer=EllipsisExpr() if self.has_default else None, # Only used by stubgen
kind=arg_kind,
)

def expand_type(self, current_info: TypeInfo) -> Type | None:
if self.type is not None and self.info.self_type is not None:
# In general, it is not safe to call `expand_type()` during semantic analyzis,
# In general, it is not safe to call `expand_type()` during semantic analysis,
# however this plugin is called very late, so all types should be fully ready.
# Also, it is tricky to avoid eager expansion of Self types here (e.g. because
# we serialize attributes).
Expand Down Expand Up @@ -269,11 +270,17 @@ def transform(self) -> bool:
if arg.kind == ARG_POS:
arg.kind = ARG_OPT

nameless_var = Var("")
existing_args_names = {arg.variable.name for arg in args}
gen_args_name = "generated_args"
while gen_args_name in existing_args_names:
gen_args_name += "_"
gen_kwargs_name = "generated_kwargs"
while gen_kwargs_name in existing_args_names:
gen_kwargs_name += "_"
args = [
Argument(nameless_var, AnyType(TypeOfAny.explicit), None, ARG_STAR),
Argument(Var(gen_args_name), AnyType(TypeOfAny.explicit), None, ARG_STAR),
*args,
Argument(nameless_var, AnyType(TypeOfAny.explicit), None, ARG_STAR2),
Argument(Var(gen_kwargs_name), AnyType(TypeOfAny.explicit), None, ARG_STAR2),
]

add_method_to_class(
Expand Down
11 changes: 0 additions & 11 deletions mypy/stubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -537,17 +537,6 @@ def _get_func_args(self, o: FuncDef, ctx: FunctionContext) -> list[ArgSig]:
if new_args is not None:
args = new_args

is_dataclass_generated = (
self.analyzed and self.processing_dataclass and o.info.names[o.name].plugin_generated
)
if o.name == "__init__" and is_dataclass_generated and "**" in [a.name for a in args]:
# The dataclass plugin generates invalid nameless "*" and "**" arguments
new_name = "".join(a.name.strip("*") for a in args)
for arg in args:
if arg.name == "*":
arg.name = f"*{new_name}_" # this name is guaranteed to be unique
elif arg.name == "**":
arg.name = f"**{new_name}__" # same here
return args

def _get_func_return(self, o: FuncDef, ctx: FunctionContext) -> str | None:
Expand Down
6 changes: 6 additions & 0 deletions test-data/unit/check-dataclasses.test
Original file line number Diff line number Diff line change
Expand Up @@ -1610,10 +1610,16 @@ B: Any
@dataclass
class A(B):
a: int
@dataclass
class C(B):
generated_args: int
generated_kwargs: int

A(a=1, b=2)
A(1)
A(a="foo") # E: Argument "a" to "A" has incompatible type "str"; expected "int"
C(generated_args="foo", generated_kwargs="bar") # E: Argument "generated_args" to "C" has incompatible type "str"; expected "int" \
# E: Argument "generated_kwargs" to "C" has incompatible type "str"; expected "int"
[builtins fixtures/dataclasses.pyi]

[case testDataclassesCallableFrozen]
Expand Down
47 changes: 32 additions & 15 deletions test-data/unit/stubgen.test
Original file line number Diff line number Diff line change
Expand Up @@ -4083,20 +4083,21 @@ class W: ...
class V: ...

[case testDataclass_semanal]
from dataclasses import dataclass, InitVar
from dataclasses import InitVar, dataclass, field
from typing import ClassVar

@dataclass
class X:
a: int
b: str = "hello"
c: ClassVar
d: ClassVar = 200
b: InitVar[str]
c: str = "hello"
d: ClassVar
e: ClassVar = 200
f: list[int] = field(init=False, default_factory=list)
g: int = field(default=2, kw_only=True)
h: int = 1
i: InitVar[str]
j: InitVar = 100
i: InitVar = 100
j: list[int] = field(default_factory=list)
non_field = None

@dataclass(init=False, repr=False, frozen=True)
Expand All @@ -4109,23 +4110,24 @@ from typing import ClassVar
@dataclass
class X:
a: int
b: str = ...
c: ClassVar
d: ClassVar = ...
b: InitVar[str]
c: str = ...
d: ClassVar
e: ClassVar = ...
f: list[int] = ...
g: int = ...
h: int = ...
i: InitVar[str]
j: InitVar = ...
i: InitVar = ...
j: list[int] = ...
non_field = ...
def __init__(self, a, b, f, g, h, i, j) -> None: ...
def __init__(self, a, b, c=..., *, g=..., h=..., i=..., j=...) -> None: ...

@dataclass(init=False, repr=False, frozen=True)
class Y: ...

[case testDataclassWithKwOnlyField_semanal]
# flags: --python-version=3.10
from dataclasses import dataclass, InitVar, KW_ONLY
from dataclasses import dataclass, field, InitVar, KW_ONLY
from typing import ClassVar

@dataclass
Expand Down Expand Up @@ -4162,7 +4164,7 @@ class X:
i: InitVar[str]
j: InitVar = ...
non_field = ...
def __init__(self, a, b, f, g, *, h, i, j) -> None: ...
def __init__(self, a, b=..., *, g=..., h=..., i, j=...) -> None: ...

@dataclass(init=False, repr=False, frozen=True)
class Y: ...
Expand Down Expand Up @@ -4193,14 +4195,29 @@ import missing
class X(missing.Base):
a: int

@dataclass
class Y(missing.Base):
generated_args: str
generated_args_: str
generated_kwargs: float
generated_kwargs_: float

[out]
import missing
from dataclasses import dataclass

@dataclass
class X(missing.Base):
a: int
def __init__(self, *selfa_, a, **selfa__) -> None: ...
def __init__(self, *generated_args, a, **generated_kwargs) -> None: ...

@dataclass
class Y(missing.Base):
generated_args: str
generated_args_: str
generated_kwargs: float
generated_kwargs_: float
def __init__(self, *generated_args__, generated_args, generated_args_, generated_kwargs, generated_kwargs_, **generated_kwargs__) -> None: ...

[case testAlwaysUsePEP604Union]
import typing
Expand Down

0 comments on commit bfbac5e

Please sign in to comment.