Skip to content

Commit

Permalink
[stubgen] Add required ... rhs to NamedTuple fields with default …
Browse files Browse the repository at this point in the history
…values (#15680)

Closes #15638
  • Loading branch information
sobolevn authored Jul 28, 2023
1 parent b901d21 commit a8467c4
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 1 deletion.
19 changes: 18 additions & 1 deletion mypy/stubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@
OverloadedFuncDef,
Statement,
StrExpr,
TempNode,
TupleExpr,
TypeInfo,
UnaryExpr,
Expand Down Expand Up @@ -637,6 +638,7 @@ def __init__(
self._state = EMPTY
self._toplevel_names: list[str] = []
self._include_private = include_private
self._current_class: ClassDef | None = None
self.import_tracker = ImportTracker()
# Was the tree semantically analysed before?
self.analyzed = analyzed
Expand Down Expand Up @@ -886,6 +888,7 @@ def get_fullname(self, expr: Expression) -> str:
return resolved_name

def visit_class_def(self, o: ClassDef) -> None:
self._current_class = o
self.method_names = find_method_names(o.defs.body)
sep: int | None = None
if not self._indent and self._state != EMPTY:
Expand Down Expand Up @@ -922,6 +925,7 @@ def visit_class_def(self, o: ClassDef) -> None:
else:
self._state = CLASS
self.method_names = set()
self._current_class = None

def get_base_types(self, cdef: ClassDef) -> list[str]:
"""Get list of base classes for a class."""
Expand Down Expand Up @@ -1330,7 +1334,20 @@ def get_init(
typename += f"[{final_arg}]"
else:
typename = self.get_str_type_of_node(rvalue)
return f"{self._indent}{lvalue}: {typename}\n"
initializer = self.get_assign_initializer(rvalue)
return f"{self._indent}{lvalue}: {typename}{initializer}\n"

def get_assign_initializer(self, rvalue: Expression) -> str:
"""Does this rvalue need some special initializer value?"""
if self._current_class and self._current_class.info:
# Current rules
# 1. Return `...` if we are dealing with `NamedTuple` and it has an existing default value
if self._current_class.info.is_named_tuple and not isinstance(rvalue, TempNode):
return " = ..."
# TODO: support other possible cases, where initializer is important

# By default, no initializer is required:
return ""

def add(self, string: str) -> None:
"""Add text to generated stub."""
Expand Down
56 changes: 56 additions & 0 deletions test-data/unit/stubgen.test
Original file line number Diff line number Diff line change
Expand Up @@ -698,6 +698,62 @@ class Y(NamedTuple):
a: int
b: str

[case testNamedTupleClassSyntax_semanal]
from typing import NamedTuple

class A(NamedTuple):
x: int
y: str = 'a'

class B(A):
z1: str
z2 = 1
z3: str = 'b'

class RegularClass:
x: int
y: str = 'a'
class NestedNamedTuple(NamedTuple):
x: int
y: str = 'a'
z: str = 'b'
[out]
from typing import NamedTuple

class A(NamedTuple):
x: int
y: str = ...

class B(A):
z1: str
z2: int
z3: str

class RegularClass:
x: int
y: str
class NestedNamedTuple(NamedTuple):
x: int
y: str = ...
z: str


[case testNestedClassInNamedTuple_semanal-xfail]
from typing import NamedTuple

# TODO: make sure that nested classes in `NamedTuple` are supported:
class NamedTupleWithNestedClass(NamedTuple):
class Nested:
x: int
y: str = 'a'
[out]
from typing import NamedTuple

class NamedTupleWithNestedClass(NamedTuple):
class Nested:
x: int
y: str

[case testEmptyNamedtuple]
import collections, typing
X = collections.namedtuple('X', [])
Expand Down

0 comments on commit a8467c4

Please sign in to comment.