diff --git a/mypy/stubgen.py b/mypy/stubgen.py index 543fbba6bf59..6a5c832704cd 100755 --- a/mypy/stubgen.py +++ b/mypy/stubgen.py @@ -43,6 +43,7 @@ import argparse import glob +import keyword import os import os.path import sys @@ -80,6 +81,7 @@ ClassDef, ComparisonExpr, Decorator, + DictExpr, EllipsisExpr, Expression, FloatExpr, @@ -126,6 +128,7 @@ from mypy.traverser import all_yield_expressions, has_return_statement, has_yield_expression from mypy.types import ( OVERLOAD_NAMES, + TPDICT_NAMES, AnyType, CallableType, Instance, @@ -405,6 +408,14 @@ def visit_tuple_expr(self, node: TupleExpr) -> str: def visit_list_expr(self, node: ListExpr) -> str: return f"[{', '.join(n.accept(self) for n in node.items)}]" + def visit_dict_expr(self, o: DictExpr) -> str: + dict_items = [] + for key, value in o.items: + # This is currently only used for TypedDict where all keys are strings. + assert isinstance(key, StrExpr) + dict_items.append(f"{key.accept(self)}: {value.accept(self)}") + return f"{{{', '.join(dict_items)}}}" + def visit_ellipsis(self, node: EllipsisExpr) -> str: return "..." @@ -641,6 +652,7 @@ def visit_mypy_file(self, o: MypyFile) -> None: "_typeshed": ["Incomplete"], "typing": ["Any", "TypeVar"], "collections.abc": ["Generator"], + "typing_extensions": ["TypedDict"], } for pkg, imports in known_imports.items(): for t in imports: @@ -1003,6 +1015,13 @@ def visit_assignment_stmt(self, o: AssignmentStmt) -> None: assert isinstance(o.rvalue, CallExpr) self.process_namedtuple(lvalue, o.rvalue) continue + if ( + isinstance(lvalue, NameExpr) + and isinstance(o.rvalue, CallExpr) + and self.is_typeddict(o.rvalue) + ): + self.process_typeddict(lvalue, o.rvalue) + continue if ( isinstance(lvalue, NameExpr) and not self.is_private_name(lvalue.name) @@ -1071,6 +1090,75 @@ def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None: self.add(f"{self._indent} {item}: Incomplete\n") self._state = CLASS + def is_typeddict(self, expr: CallExpr) -> bool: + callee = expr.callee + return ( + isinstance(callee, NameExpr) and self.refers_to_fullname(callee.name, TPDICT_NAMES) + ) or ( + isinstance(callee, MemberExpr) + and isinstance(callee.expr, NameExpr) + and f"{callee.expr.name}.{callee.name}" in TPDICT_NAMES + ) + + def process_typeddict(self, lvalue: NameExpr, rvalue: CallExpr) -> None: + if self._state != EMPTY: + self.add("\n") + + if not isinstance(rvalue.args[0], StrExpr): + self.add(f"{self._indent}{lvalue.name}: Incomplete") + self.import_tracker.require_name("Incomplete") + return + + items: list[tuple[str, Expression]] = [] + total: Expression | None = None + if len(rvalue.args) > 1 and rvalue.arg_kinds[1] == ARG_POS: + if not isinstance(rvalue.args[1], DictExpr): + self.add(f"{self._indent}{lvalue.name}: Incomplete") + self.import_tracker.require_name("Incomplete") + return + for attr_name, attr_type in rvalue.args[1].items: + if not isinstance(attr_name, StrExpr): + self.add(f"{self._indent}{lvalue.name}: Incomplete") + self.import_tracker.require_name("Incomplete") + return + items.append((attr_name.value, attr_type)) + if len(rvalue.args) > 2: + if rvalue.arg_kinds[2] != ARG_NAMED or rvalue.arg_names[2] != "total": + self.add(f"{self._indent}{lvalue.name}: Incomplete") + self.import_tracker.require_name("Incomplete") + return + total = rvalue.args[2] + else: + for arg_name, arg in zip(rvalue.arg_names[1:], rvalue.args[1:]): + if not isinstance(arg_name, str): + self.add(f"{self._indent}{lvalue.name}: Incomplete") + self.import_tracker.require_name("Incomplete") + return + if arg_name == "total": + total = arg + else: + items.append((arg_name, arg)) + self.import_tracker.require_name("TypedDict") + p = AliasPrinter(self) + if any(not key.isidentifier() or keyword.iskeyword(key) for key, _ in items): + # Keep the call syntax if there are non-identifier or keyword keys. + self.add(f"{self._indent}{lvalue.name} = {rvalue.accept(p)}\n") + self._state = VAR + else: + bases = "TypedDict" + # TODO: Add support for generic TypedDicts. Requires `Generic` as base class. + if total is not None: + bases += f", total={total.accept(p)}" + self.add(f"{self._indent}class {lvalue.name}({bases}):") + if len(items) == 0: + self.add(" ...\n") + self._state = EMPTY_CLASS + else: + self.add("\n") + for key, key_type in items: + self.add(f"{self._indent} {key}: {key_type.accept(p)}\n") + self._state = CLASS + def is_alias_expression(self, expr: Expression, top_level: bool = True) -> bool: """Return True for things that look like target for an alias. diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test index 8e4285b7de2e..9b3b8e3209b5 100644 --- a/test-data/unit/stubgen.test +++ b/test-data/unit/stubgen.test @@ -2793,3 +2793,116 @@ def f(x: str | None) -> None: ... a: str | int def f(x: str | None) -> None: ... + +[case testTypeddict] +import typing, x +X = typing.TypedDict('X', {'a': int, 'b': str}) +Y = typing.TypedDict('X', {'a': int, 'b': str}, total=False) +[out] +from typing_extensions import TypedDict + +class X(TypedDict): + a: int + b: str + +class Y(TypedDict, total=False): + a: int + b: str + +[case testTypeddictKeywordSyntax] +from typing import TypedDict + +X = TypedDict('X', a=int, b=str) +Y = TypedDict('X', a=int, b=str, total=False) +[out] +from typing import TypedDict + +class X(TypedDict): + a: int + b: str + +class Y(TypedDict, total=False): + a: int + b: str + +[case testTypeddictWithNonIdentifierOrKeywordKeys] +from typing import TypedDict +X = TypedDict('X', {'a-b': int, 'c': str}) +Y = TypedDict('X', {'a-b': int, 'c': str}, total=False) +Z = TypedDict('X', {'a': int, 'in': str}) +[out] +from typing import TypedDict + +X = TypedDict('X', {'a-b': int, 'c': str}) + +Y = TypedDict('X', {'a-b': int, 'c': str}, total=False) + +Z = TypedDict('X', {'a': int, 'in': str}) + +[case testEmptyTypeddict] +import typing +X = typing.TypedDict('X', {}) +Y = typing.TypedDict('Y', {}, total=False) +Z = typing.TypedDict('Z') +W = typing.TypedDict('W', total=False) +[out] +from typing_extensions import TypedDict + +class X(TypedDict): ... + +class Y(TypedDict, total=False): ... + +class Z(TypedDict): ... + +class W(TypedDict, total=False): ... + +[case testTypeddictAliased] +from typing import TypedDict as t_TypedDict +from typing_extensions import TypedDict as te_TypedDict +def f(): ... +X = t_TypedDict('X', {'a': int, 'b': str}) +Y = te_TypedDict('Y', {'a': int, 'b': str}) +def g(): ... +[out] +from typing_extensions import TypedDict + +def f() -> None: ... + +class X(TypedDict): + a: int + b: str + +class Y(TypedDict): + a: int + b: str + +def g() -> None: ... + +[case testNotTypeddict] +from x import TypedDict +import y +X = TypedDict('X', {'a': int, 'b': str}) +Y = y.TypedDict('Y', {'a': int, 'b': str}) +[out] +from _typeshed import Incomplete + +X: Incomplete +Y: Incomplete + +[case testTypeddictWithWrongAttributesType] +from typing import TypedDict +R = TypedDict("R", {"a": int, **{"b": str, "c": bytes}}) +S = TypedDict("S", [("b", str), ("c", bytes)]) +T = TypedDict("T", {"a": int}, b=str, total=False) +U = TypedDict("U", {"a": int}, totale=False) +V = TypedDict("V", {"a": int}, {"b": str}) +W = TypedDict("W", **{"a": int, "b": str}) +[out] +from _typeshed import Incomplete + +R: Incomplete +S: Incomplete +T: Incomplete +U: Incomplete +V: Incomplete +W: Incomplete