-
-
Notifications
You must be signed in to change notification settings - Fork 2.8k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
stubgen: Support TypedDict alternative syntax #14682
Changes from all commits
56666e6
560e744
3f2c434
0e617b6
edeba78
444b3f4
5a85e1d
1bd4c13
fe32c21
a6fbd4a
d9cc7f5
b42f9b6
f06b01c
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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") | ||
hamdanal marked this conversation as resolved.
Show resolved
Hide resolved
|
||
self._state = VAR | ||
else: | ||
bases = "TypedDict" | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Mypy actually now supports generic TypedDicts defined with the call syntax, in which case I think it's reasonable not to handle generic TypedDicts defined using the call syntax in this PR, since they're unlikely to come up much. But probably worth a TODO comment? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. TIL, thank you. |
||
# 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. | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I still think it would be good to be more cautious here. What if somebody runs stubgen on code like this?
Now, you'll tell me that this code raises
TypeError
at runtime, and indeed it does! But that doesn't mean that stubgen should crash (or do something incorrect) if it's run on code like this. We should just emitT: Incomplete
and move on to the next thing in the file.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I added more checks for invalid cases including this one. I also added a test for importing from
typing_extensions
.With all the checks for invalid uses, the
TypedDict
handling part looks much more conservative than other parts of stubgen. I am not sure what to make of this, it is not necessarily a positive or negative thing, just something I noticed.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hmm, well, I'm not really an expert on stubgen, so I can't really comment on decisions that have previously been taken. But in my opinion, it's always good for static analysis tools to be as paranoid as possible about the kind of thing that might be fed to them :)