diff --git a/mypy/stubgen.py b/mypy/stubgen.py index 91f461b84c15d..1959f71d12ed4 100755 --- a/mypy/stubgen.py +++ b/mypy/stubgen.py @@ -91,7 +91,7 @@ from mypy.find_sources import create_source_list, InvalidSourceList from mypy.build import build from mypy.errors import CompileError, Errors -from mypy.traverser import has_return_statement +from mypy.traverser import all_yield_expressions, has_return_statement, has_yield_expression from mypy.moduleinspect import ModuleInspect @@ -550,13 +550,17 @@ def visit_mypy_file(self, o: MypyFile) -> None: self.path = o.path self.defined_names = find_defined_names(o) self.referenced_names = find_referenced_names(o) - typing_imports = ["Any", "Optional", "TypeVar"] - for t in typing_imports: - if t not in self.defined_names: - alias = None - else: - alias = '_' + t - self.import_tracker.add_import_from("typing", [(t, alias)]) + known_imports = { + "typing": ["Any", "TypeVar"], + "collections.abc": ["Generator"], + } + for pkg, imports in known_imports.items(): + for t in imports: + if t not in self.defined_names: + alias = None + else: + alias = '_' + t + self.import_tracker.add_import_from(pkg, [(t, alias)]) super().visit_mypy_file(o) undefined_names = [name for name in self._all_ or [] if name not in self._toplevel_names] @@ -662,6 +666,23 @@ def visit_func_def(self, o: FuncDef, is_abstract: bool = False, # Always assume abstract methods return Any unless explicitly annotated. Also # some dunder methods should not have a None return type. retname = None # implicit Any + elif has_yield_expression(o): + self.add_abc_import('Generator') + yield_name = 'None' + send_name = 'None' + return_name = 'None' + 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('Any') + yield_name = 'Any' + if in_assignment: + self.add_typing_import('Any') + send_name = 'Any' + if has_return_statement(o): + self.add_typing_import('Any') + return_name = 'Any' + generator_name = self.typing_name('Generator') + retname = f'{generator_name}[{yield_name}, {send_name}, {return_name}]' elif not has_return_statement(o) and not is_abstract: retname = 'None' retfield = '' @@ -672,6 +693,9 @@ def visit_func_def(self, o: FuncDef, is_abstract: bool = False, self.add("){}: ...\n".format(retfield)) self._state = FUNC + def is_none_expr(self, expr: Expression) -> bool: + return isinstance(expr, NameExpr) and expr.name == "None" + def visit_decorator(self, o: Decorator) -> None: if self.is_private_name(o.func.name, o.func.fullname): return @@ -1107,6 +1131,14 @@ def add_typing_import(self, name: str) -> None: name = self.typing_name(name) self.import_tracker.require_name(name) + def add_abc_import(self, name: str) -> None: + """Add a name to be imported from collections.abc, unless it's imported already. + + The import will be internal to the stub. + """ + name = self.typing_name(name) + self.import_tracker.require_name(name) + def add_import_line(self, line: str) -> None: """Add a line of text to the import section, unless it's already there.""" if line not in self._import_lines: diff --git a/mypy/traverser.py b/mypy/traverser.py index d0b656c7a77f2..a5f993bd2fa52 100644 --- a/mypy/traverser.py +++ b/mypy/traverser.py @@ -1,6 +1,6 @@ """Generic node traverser visitor""" -from typing import List +from typing import List, Tuple from mypy_extensions import mypyc_attr from mypy.visitor import NodeVisitor @@ -319,9 +319,22 @@ def has_return_statement(fdef: FuncBase) -> bool: return seeker.found -class ReturnCollector(TraverserVisitor): +class YieldSeeker(TraverserVisitor): + def __init__(self) -> None: + self.found = False + + def visit_yield_expr(self, o: YieldExpr) -> None: + self.found = True + + +def has_yield_expression(fdef: FuncBase) -> bool: + seeker = YieldSeeker() + fdef.accept(seeker) + return seeker.found + + +class FuncCollectorBase(TraverserVisitor): def __init__(self) -> None: - self.return_statements: List[ReturnStmt] = [] self.inside_func = False def visit_func_def(self, defn: FuncDef) -> None: @@ -330,6 +343,12 @@ def visit_func_def(self, defn: FuncDef) -> None: super().visit_func_def(defn) self.inside_func = False + +class ReturnCollector(FuncCollectorBase): + def __init__(self) -> None: + super().__init__() + self.return_statements: List[ReturnStmt] = [] + def visit_return_stmt(self, stmt: ReturnStmt) -> None: self.return_statements.append(stmt) @@ -338,3 +357,24 @@ def all_return_statements(node: Node) -> List[ReturnStmt]: v = ReturnCollector() node.accept(v) return v.return_statements + + +class YieldCollector(FuncCollectorBase): + def __init__(self) -> None: + super().__init__() + self.in_assignment = False + self.yield_expressions: List[Tuple[YieldExpr, bool]] = [] + + def visit_assignment_stmt(self, stmt: AssignmentStmt) -> None: + self.in_assignment = True + super().visit_assignment_stmt(stmt) + self.in_assignment = False + + def visit_yield_expr(self, expr: YieldExpr) -> None: + self.yield_expressions.append((expr, self.in_assignment)) + + +def all_yield_expressions(node: Node) -> List[Tuple[YieldExpr, bool]]: + v = YieldCollector() + node.accept(v) + return v.yield_expressions diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test index daee45a840825..b1995b57e7759 100644 --- a/test-data/unit/stubgen.test +++ b/test-data/unit/stubgen.test @@ -944,6 +944,56 @@ def f(): ... [out] def f() -> None: ... +[case testFunctionYields] +def f(): + yield 123 +def g(): + x = yield +def h1(): + yield + return +def h2(): + yield + return "abc" +def all(): + x = yield 123 + return "abc" +[out] +from collections.abc import Generator +from typing import Any + +def f() -> Generator[Any, None, None]: ... +def g() -> Generator[None, Any, None]: ... +def h1() -> Generator[None, None, None]: ... +def h2() -> Generator[None, None, Any]: ... +def all() -> Generator[Any, Any, Any]: ... + +[case testFunctionYieldsNone] +def f(): + yield +def g(): + yield None + +[out] +from collections.abc import Generator + +def f() -> Generator[None, None, None]: ... +def g() -> Generator[None, None, None]: ... + +[case testGeneratorAlreadyDefined] +class Generator: + pass + +def f(): + yield 123 +[out] +from collections.abc import Generator as _Generator +from typing import Any + +class Generator: ... + +def f() -> _Generator[Any, None, None]: ... + [case testCallable] from typing import Callable