Skip to content
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.

Sign up for GitHub

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

(🎁) StubGenerator.add_typing_import returns the name #15912

Merged
merged 1 commit into from
Aug 21, 2023
Merged

(🎁) StubGenerator.add_typing_import returns the name #15912

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 22 additions & 33 deletions mypy/stubgen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -786,25 +786,20 @@ def visit_func_def(self, o: FuncDef) -> None:
elif o.name in KNOWN_MAGIC_METHODS_RETURN_TYPES:
retname = KNOWN_MAGIC_METHODS_RETURN_TYPES[o.name]
elif has_yield_expression(o) or has_yield_from_expression(o):
self.add_typing_import("Generator")
generator_name = self.add_typing_import("Generator")
yield_name = "None"
send_name = "None"
return_name = "None"
if has_yield_from_expression(o):
self.add_typing_import("Incomplete")
yield_name = send_name = self.typing_name("Incomplete")
yield_name = send_name = self.add_typing_import("Incomplete")
else:
for expr, in_assignment in all_yield_expressions(o):
if expr.expr is not None and not self.is_none_expr(expr.expr):
self.add_typing_import("Incomplete")
yield_name = self.typing_name("Incomplete")
yield_name = self.add_typing_import("Incomplete")
if in_assignment:
self.add_typing_import("Incomplete")
send_name = self.typing_name("Incomplete")
send_name = self.add_typing_import("Incomplete")
if has_return_statement(o):
self.add_typing_import("Incomplete")
return_name = self.typing_name("Incomplete")
generator_name = self.typing_name("Generator")
return_name = self.add_typing_import("Incomplete")
retname = f"{generator_name}[{yield_name}, {send_name}, {return_name}]"
elif not has_return_statement(o) and o.abstract_status == NOT_ABSTRACT:
retname = "None"
Expand Down Expand Up @@ -965,21 +960,19 @@ def get_base_types(self, cdef: ClassDef) -> list[str]:
nt_fields = self._get_namedtuple_fields(base)
assert isinstance(base.args[0], StrExpr)
typename = base.args[0].value
if nt_fields is not None:
fields_str = ", ".join(f"({f!r}, {t})" for f, t in nt_fields)
namedtuple_name = self.typing_name("NamedTuple")
base_types.append(f"{namedtuple_name}({typename!r}, [{fields_str}])")
self.add_typing_import("NamedTuple")
else:
if nt_fields is None:
# Invalid namedtuple() call, cannot determine fields
base_types.append(self.typing_name("Incomplete"))
base_types.append(self.add_typing_import("Incomplete"))
continue
fields_str = ", ".join(f"({f!r}, {t})" for f, t in nt_fields)
namedtuple_name = self.add_typing_import("NamedTuple")
base_types.append(f"{namedtuple_name}({typename!r}, [{fields_str}])")
elif self.is_typed_namedtuple(base):
base_types.append(base.accept(p))
else:
# At this point, we don't know what the base class is, so we
# just use Incomplete as the base class.
base_types.append(self.typing_name("Incomplete"))
self.add_typing_import("Incomplete")
base_types.append(self.add_typing_import("Incomplete"))
for name, value in cdef.keywords.items():
if name == "metaclass":
continue # handled separately
Expand Down Expand Up @@ -1059,9 +1052,9 @@ def _get_namedtuple_fields(self, call: CallExpr) -> list[tuple[str, str]] | None
field_names.append(field.value)
else:
return None # Invalid namedtuple fields type
if field_names:
self.add_typing_import("Incomplete")
incomplete = self.typing_name("Incomplete")
if not field_names:
return []
incomplete = self.add_typing_import("Incomplete")
return [(field_name, incomplete) for field_name in field_names]
elif self.is_typed_namedtuple(call):
fields_arg = call.args[1]
Expand Down Expand Up @@ -1092,8 +1085,7 @@ def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
if fields is None:
self.annotate_as_incomplete(lvalue)
return
self.add_typing_import("NamedTuple")
bases = self.typing_name("NamedTuple")
bases = self.add_typing_import("NamedTuple")
# TODO: Add support for generic NamedTuples. Requires `Generic` as base class.
class_def = f"{self._indent}class {lvalue.name}({bases}):"
if len(fields) == 0:
Expand Down Expand Up @@ -1143,14 +1135,13 @@ def process_typeddict(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
total = arg
else:
items.append((arg_name, arg))
self.add_typing_import("TypedDict")
bases = self.add_typing_import("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 reserved keyword keys.
self.add(f"{self._indent}{lvalue.name} = {rvalue.accept(p)}\n")
self._state = VAR
else:
bases = self.typing_name("TypedDict")
# TODO: Add support for generic TypedDicts. Requires `Generic` as base class.
if total is not None:
bases += f", total={total.accept(p)}"
Expand All @@ -1167,8 +1158,7 @@ def process_typeddict(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
self._state = CLASS

def annotate_as_incomplete(self, lvalue: NameExpr) -> None:
self.add_typing_import("Incomplete")
self.add(f"{self._indent}{lvalue.name}: {self.typing_name('Incomplete')}\n")
self.add(f"{self._indent}{lvalue.name}: {self.add_typing_import('Incomplete')}\n")
self._state = VAR

def is_alias_expression(self, expr: Expression, top_level: bool = True) -> bool:
Expand Down Expand Up @@ -1384,13 +1374,14 @@ def typing_name(self, name: str) -> str:
else:
return name

def add_typing_import(self, name: str) -> None:
def add_typing_import(self, name: str) -> str:
"""Add a name to be imported for typing, unless it's imported already.

The import will be internal to the stub.
"""
name = self.typing_name(name)
self.import_tracker.require_name(name)
return name

def add_import_line(self, line: str) -> None:
"""Add a line of text to the import section, unless it's already there."""
Expand Down Expand Up @@ -1448,11 +1439,9 @@ def get_str_type_of_node(
if isinstance(rvalue, NameExpr) and rvalue.name in ("True", "False"):
return "bool"
if can_infer_optional and isinstance(rvalue, NameExpr) and rvalue.name == "None":
self.add_typing_import("Incomplete")
return f"{self.typing_name('Incomplete')} | None"
return f"{self.add_typing_import('Incomplete')} | None"
if can_be_any:
self.add_typing_import("Incomplete")
return self.typing_name("Incomplete")
return self.add_typing_import("Incomplete")
else:
return ""

Expand Down