10000 stubgen: Use NamedTuple class syntax (#10625) · python/mypy@fee4422 · GitHub
[go: up one dir, main page]
Skip to content

Commit fee4422

Browse files
srittauJukkaL
srittau
authored and
JukkaL
committed
stubgen: Use NamedTuple class syntax (#10625)
### Description Use `typing.NamedTuple` in stubgen, instead of `collections.namedtuple`. ## Test Plan I modified the existing tests and added another test for an empty namedtuple, to ensure that `Any` is not imported unnecessarily.
1 parent 5af694e commit fee4422

File tree

2 files changed

+50
-20
lines changed

2 files changed

+50
-20
lines changed

‎mypy/stubgen.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -536,7 +536,7 @@ def __init__(self,
536536
# Disable implicit exports of package-internal imports?
537537
self.export_less = export_less
538538
# Add imports that could be implicitly generated
539-
self.import_tracker.add_import_from("collections", [("namedtuple", None)])
539+
self.import_tracker.add_import_from("typing", [("NamedTuple", None)])
540540
# Names in __all__ are required
541541
for name in _all_ or ():
542542
if name not in IGNORED_DUNDERS:
@@ -900,18 +900,24 @@ def is_namedtuple(self, expr: Expression) -> bool:
900900
def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
901901
if self._state != EMPTY:
902902
self.add('\n')
903-
name = repr(getattr(rvalue.args[0], 'value', ERROR_MARKER))
904903
if isinstance(rvalue.args[1], StrExpr):
905-
items = repr(rvalue.args[1].value)
904+
items = rvalue.args[1].value.split(" ")
906905
elif isinstance(rvalue.args[1], (ListExpr, TupleExpr)):
907906
list_items = cast(List[StrExpr], rvalue.args[1].items)
908-
items = '[%s]' % ', '.join(repr(item.value) for item in list_items)
907+
items = [item.value for item in list_items]
909908
else:
910909
self.add('%s%s: Any' % (self._indent, lvalue.name))
911910
self.import_tracker.require_name('Any')
912911
return
913-
self.import_tracker.require_name('namedtuple')
914-
self.add('%s%s = namedtuple(%s, %s)\n' % (self._indent, lvalue.name, name, items))
912+
self.import_tracker.require_name('NamedTuple')
913+
self.add('{}class {}(NamedTuple):'.format(self._indent, lvalue.name))
914+
if len(items) == 0:
915+
self.add(' ...\n')
916+
else:
917+
self.import_tracker.require_name('Any')
918+
self.add('\n')
919+
for item in items:
920+
self.add('{} {}: Any\n'.format(self._indent, item))
915921
self._state = CLASS
916922

917923
def is_alias_expression(self, expr: Expression, top_level: bool = True) -> bool:

‎test-data/unit/stubgen.test

Lines changed: 38 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -591,30 +591,44 @@ class A:
591591
import collections, x
592592
X = collections.namedtuple('X', ['a', 'b'])
593593
[out]
594-
from collections import namedtuple
594+
from typing import Any, NamedTuple
595+
596+
class X(NamedTuple):
597+
a: Any
598+
b: Any
595599

596-
X = namedtuple('X', ['a', 'b'])
600+
[case testEmptyNamedtuple]
601+
import collections
602+
X = collections.namedtuple('X', [])
603+
[out]
604+
from typing import NamedTuple
605+
606+
class X(NamedTuple): ...
597607

598608
[case testNamedtupleAltSyntax]
599609
from collections import namedtuple, xx
600610
X = namedtuple('X', 'a b')
601611
xx
602612
[out]
603-
from collections import namedtuple
613+
from typing import Any, NamedTuple
604614

605-
X = namedtuple('X', 'a b')
615+
class X(NamedTuple):
616+
a: Any
617+
b: Any
606618

607619
[case testNamedtupleWithUnderscore]
608620
from collections import namedtuple as _namedtuple
609621
def f(): ...
610622
X = _namedtuple('X', 'a b')
611623
def g(): ...
612624
[out]
613-
from collections import namedtuple
625+
from typing import Any, NamedTuple
614626

615627
def f() -> None: ...
616628

617-
X = namedtuple('X', 'a b')
629+
class X(NamedTuple):
630+
a: Any
631+
b: Any
618632

619633
def g() -> None: ...
620634

@@ -623,9 +637,11 @@ import collections, x
623637
_X = collections.namedtuple('_X', ['a', 'b'])
624638
class Y(_X): ...
625639
[out]
626-
from collections import namedtuple
640+
from typing import Any, NamedTuple
627641

628-
_X = namedtuple('_X', ['a', 'b'])
642+
class _X(NamedTuple):
643+
a: Any
644+
b: Any
629645

630646
class Y(_X): ...
631647

@@ -636,13 +652,19 @@ Y = namedtuple('Y', ('a',))
636652
Z = namedtuple('Z', ('a', 'b', 'c', 'd', 'e'))
637653
xx
638654
[out]
639-
from collections import namedtuple
655+
from typing import Any, NamedTuple
640656

641-
X = namedtuple('X', [])
657+
class X(NamedTuple): ...
642658

643-
Y = namedtuple('Y', ['a'])
659+
class Y(NamedTuple):
660+
a: Any
644661

645-
Z = namedtuple('Z', ['a', 'b', 'c', 'd', 'e'])
662+
class Z(NamedTuple):
663+
a: Any
664+
b: Any
665+
c: Any
666+
d: Any
667+
e: Any
646668

647669
[case testDynamicNamedTuple]
648670
from collections import namedtuple
@@ -2187,10 +2209,12 @@ from collections import namedtuple
21872209
class C:
21882210
N = namedtuple('N', ['x', 'y'])
21892211
[out]
2190-
from collections import namedtuple
2212+
from typing import Any, NamedTuple
21912213

21922214
class C:
2193-
N = namedtuple('N', ['x', 'y'])
2215+
class N(NamedTuple):
2216+
x: Any
2217+
y: Any
21942218

21952219
[case testImports_directImportsWithAlias]
21962220
import p.a as a

0 commit comments

Comments
 (0)
0