Support for typechecking generators (#529)

JelleZijlstra · web-flow · commit 9446c3571383 · 2022-04-21T12:42:57.000-07:00
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Type check `yield`, `yield from`, and `return` nodes in generators (#529)
 - Type check calls to comparison operators (#527)
 - Retrieve attributes from stubs even when a runtime
   equivalent exists (#526)
diff --git a/pyanalyze/annotations.py b/pyanalyze/annotations.py
@@ -76,6 +76,7 @@
 )
 from .value import (
     _HashableValue,
+    TypeVarLike,
     annotate_value,
     AnnotatedValue,
     AnySource,
@@ -610,12 +611,12 @@ def _type_from_runtime(
         return AnyValue(AnySource.error)
 
 
-def make_type_var_value(tv: TypeVar, ctx: Context) -> TypeVarValue:
+def make_type_var_value(tv: TypeVarLike, ctx: Context) -> TypeVarValue:
     if tv.__bound__ is not None:
         bound = _type_from_runtime(tv.__bound__, ctx)
     else:
         bound = None
-    if tv.__constraints__:
+    if isinstance(tv, TypeVar) and tv.__constraints__:
         constraints = tuple(
             _type_from_runtime(constraint, ctx) for constraint in tv.__constraints__
         )
diff --git a/pyanalyze/error_code.py b/pyanalyze/error_code.py
@@ -98,6 +98,7 @@ class ErrorCode(enum.Enum):
     already_declared = 78
     invalid_annotated_assignment = 79
     unused_assignment = 80
+    incompatible_yield = 81
 
 
 # Allow testing unannotated functions without too much fuss
@@ -217,6 +218,7 @@ class ErrorCode(enum.Enum):
     ErrorCode.already_declared: "Name is already declared",
     ErrorCode.invalid_annotated_assignment: "Invalid annotated assignment",
     ErrorCode.unused_assignment: "Assigned value is never used",
+    ErrorCode.incompatible_yield: "Incompatible yield type",
 }
 
 
diff --git a/pyanalyze/functions.py b/pyanalyze/functions.py
@@ -23,13 +23,15 @@
 from .options import Options, PyObjectSequenceOption
 from .signature import ParameterKind, Signature, SigParameter
 from .stacked_scopes import Composite
+from .typevar import resolve_bounds_map
 from .value import (
     AnySource,
     AnyValue,
     CallableValue,
     CanAssignContext,
     CanAssignError,
     GenericValue,
+    get_tv_map,
     KnownValue,
     SubclassValue,
     TypedValue,
@@ -43,6 +45,15 @@
 FunctionNode = Union[FunctionDefNode, ast.Lambda]
 IMPLICIT_CLASSMETHODS = ("__init_subclass__", "__new__")
 
+YieldT = TypeVar("YieldT")
+SendT = TypeVar("SendT")
+ReturnT = TypeVar("ReturnT")
+IterableValue = GenericValue(collections.abc.Iterable, [TypeVarValue(YieldT)])
+GeneratorValue = GenericValue(
+    collections.abc.Generator,
+    [TypeVarValue(YieldT), TypeVarValue(SendT), TypeVarValue(ReturnT)],
+)
+
 
 class AsyncFunctionKind(enum.Enum):
     non_async = 0
@@ -78,6 +89,41 @@ class FunctionInfo:
     return_annotation: Optional[Value]
     potential_function: Optional[object]
 
+    def get_generator_yield_type(self, ctx: CanAssignContext) -> Value:
+        if self.return_annotation is None:
+            return AnyValue(AnySource.unannotated)
+        can_assign = IterableValue.can_assign(self.return_annotation, ctx)
+        if isinstance(can_assign, CanAssignError):
+            return AnyValue(AnySource.error)
+        tv_map, _ = resolve_bounds_map(can_assign, ctx)
+        return tv_map.get(YieldT, AnyValue(AnySource.generic_argument))
+
+    def get_generator_send_type(self, ctx: CanAssignContext) -> Value:
+        if self.return_annotation is None:
+            return AnyValue(AnySource.unannotated)
+        tv_map = get_tv_map(GeneratorValue, self.return_annotation, ctx)
+        if not isinstance(tv_map, CanAssignError):
+            return tv_map.get(SendT, AnyValue(AnySource.generic_argument))
+        # If the return annotation is a non-Generator Iterable, assume the send
+        # type is None.
+        can_assign = IterableValue.can_assign(self.return_annotation, ctx)
+        if isinstance(can_assign, CanAssignError):
+            return AnyValue(AnySource.error)
+        return KnownValue(None)
+
+    def get_generator_return_type(self, ctx: CanAssignContext) -> Value:
+        if self.return_annotation is None:
+            return AnyValue(AnySource.unannotated)
+        tv_map = get_tv_map(GeneratorValue, self.return_annotation, ctx)
+        if not isinstance(tv_map, CanAssignError):
+            return tv_map.get(ReturnT, AnyValue(AnySource.generic_argument))
+        # If the return annotation is a non-Generator Iterable, assume the return
+        # type is None.
+        can_assign = IterableValue.can_assign(self.return_annotation, ctx)
+        if isinstance(can_assign, CanAssignError):
+            return AnyValue(AnySource.error)
+        return KnownValue(None)
+
 
 @dataclass
 class FunctionResult:
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
@@ -82,7 +82,12 @@
     FunctionInfo,
     FunctionNode,
     FunctionResult,
+    GeneratorValue,
     IMPLICIT_CLASSMETHODS,
+    IterableValue,
+    ReturnT,
+    SendT,
+    YieldT,
 )
 from .options import (
     add_arguments,
@@ -1021,6 +1026,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
     current_class: Optional[type]
     current_enum_members: Optional[Dict[object, str]]
     current_function: Optional[object]
+    current_function_info: Optional[FunctionInfo]
     current_function_name: Optional[str]
     error_for_implicit_any: bool
     expected_return_value: Optional[Value]
@@ -1081,6 +1087,7 @@ def __init__(
         # current class (for inferring the type of cls and self arguments)
         self.current_class = None
         self.current_function_name = None
+        self.current_function_info = None
 
         # async
         self.async_kind = AsyncFunctionKind.non_async
@@ -1712,6 +1719,8 @@ def visit_FunctionDef(self, node: FunctionDefNode) -> Value:
             self, "current_function", potential_function
         ), qcore.override(
             self, "expected_return_value", expected_return
+        ), qcore.override(
+            self, "current_function_info", info
         ):
             result = self._visit_function_body(info)
 
@@ -1831,7 +1840,8 @@ def record_call(self, callable: object, arguments: CallArgs) -> None:
     def visit_Lambda(self, node: ast.Lambda) -> Value:
         with self.asynq_checker.set_func_name("<lambda>"):
             info = compute_function_info(node, self)
-            result = self._visit_function_body(info)
+            with qcore.override(self, "current_function_info", info):
+                result = self._visit_function_body(info)
             return compute_value_of_function(info, self, result=result.return_value)
 
     def _visit_function_body(self, function_info: FunctionInfo) -> FunctionResult:
@@ -3236,15 +3246,53 @@ def unpack_awaitable(self, composite: Composite, node: ast.AST) -> Value:
     def visit_YieldFrom(self, node: ast.YieldFrom) -> Value:
         self.is_generator = True
         value = self.visit(node.value)
-        if not TypedValue(collections.abc.Iterable).is_assignable(
-            value, self
-        ) and not AwaitableValue.is_assignable(value, self):
-            self._show_error_if_checking(
-                node,
-                f"Cannot use {value} in yield from",
-                error_code=ErrorCode.bad_yield_from,
-            )
-        return AnyValue(AnySource.inference)
+        tv_map = get_tv_map(GeneratorValue, value, self)
+        if isinstance(tv_map, CanAssignError):
+            can_assign = get_tv_map(AwaitableValue, value, self)
+            if not isinstance(can_assign, CanAssignError):
+                tv_map = {
+                    ReturnT: can_assign.get(T, AnyValue(AnySource.generic_argument))
+                }
+            else:
+                can_assign = get_tv_map(IterableValue, value, self)
+                if isinstance(can_assign, CanAssignError):
+                    self._show_error_if_checking(
+                        node,
+                        f"Cannot use {value} in yield from",
+                        error_code=ErrorCode.bad_yield_from,
+                        detail=can_assign.display(),
+                    )
+                    tv_map = {ReturnT: AnyValue(AnySource.error)}
+                else:
+                    tv_map = {
+                        YieldT: can_assign.get(T, AnyValue(AnySource.generic_argument))
+                    }
+
+        if self.current_function_info is not None:
+            expected_yield = self.current_function_info.get_generator_yield_type(self)
+            yield_type = tv_map.get(YieldT, AnyValue(AnySource.generic_argument))
+            can_assign = expected_yield.can_assign(yield_type, self)
+            if isinstance(can_assign, CanAssignError):
+                self._show_error_if_checking(
+                    node,
+                    f"Cannot yield from {value} (expected {expected_yield})",
+                    error_code=ErrorCode.incompatible_yield,
+                    detail=can_assign.display(),
+                )
+
+            expected_send = self.current_function_info.get_generator_send_type(self)
+            send_type = tv_map.get(SendT, AnyValue(AnySource.generic_argument))
+            can_assign = send_type.can_assign(expected_send, self)
+            if isinstance(can_assign, CanAssignError):
+                self._show_error_if_checking(
+                    node,
+                    f"Cannot send {send_type} to a generator (expected"
+                    f" {expected_send})",
+                    error_code=ErrorCode.incompatible_yield,
+                    detail=can_assign.display(),
+                )
+
+        return tv_map.get(ReturnT, AnyValue(AnySource.generic_argument))
 
     def visit_Yield(self, node: ast.Yield) -> Value:
         if self._is_checking():
@@ -3260,7 +3308,7 @@ def visit_Yield(self, node: ast.Yield) -> Value:
             if node.value is not None:
                 value = self.visit(node.value)
             else:
-                value = None
+                value = KnownValue(None)
 
         if node.value is None and self.async_kind in (
             AsyncFunctionKind.normal,
@@ -3270,10 +3318,21 @@ def visit_Yield(self, node: ast.Yield) -> Value:
         self.is_generator = True
 
         # unwrap the results of async yields
-        if self.async_kind != AsyncFunctionKind.non_async and value is not None:
+        if self.async_kind != AsyncFunctionKind.non_async:
             return self._unwrap_yield_result(node, value)
-        else:
+        if self.current_function_info is None:
             return AnyValue(AnySource.inference)
+        yield_type = self.current_function_info.get_generator_yield_type(self)
+        can_assign = yield_type.can_assign(value, self)
+        if isinstance(can_assign, CanAssignError):
+            self._show_error_if_checking(
+                node,
+                f"Cannot assign value of type {value} to yield expression of type"
+                f" {yield_type}",
+                error_code=ErrorCode.incompatible_yield,
+                detail=can_assign.display(),
+            )
+        return self.current_function_info.get_generator_send_type(self)
 
     def _unwrap_yield_result(self, node: ast.AST, value: Value) -> Value:
         if isinstance(value, AsyncTaskIncompleteValue):
@@ -3335,7 +3394,7 @@ def _unwrap_yield_result(self, node: ast.AST, value: Value) -> Value:
         else:
             self._show_error_if_checking(
                 node,
-                "Invalid value yielded: %r" % (value,),
+                f"Invalid value yielded: {value}",
                 error_code=ErrorCode.bad_async_yield,
             )
             return AnyValue(AnySource.error)
@@ -3354,11 +3413,18 @@ def visit_Return(self, node: ast.Return) -> None:
             self._show_error_if_checking(
                 node, error_code=ErrorCode.no_return_may_return
             )
-        elif (
-            # TODO check generator types properly
-            not (self.is_generator and self.async_kind == AsyncFunctionKind.non_async)
-            and self.expected_return_value is not None
-        ):
+        elif self.is_generator and self.async_kind == AsyncFunctionKind.non_async:
+            if self.current_function_info is not None:
+                expected = self.current_function_info.get_generator_return_type(self)
+                can_assign = expected.can_assign(value, self)
+                if isinstance(can_assign, CanAssignError):
+                    self._show_error_if_checking(
+                        node,
+                        f"Incompatible return value {value} (expected {expected})",
+                        error_code=ErrorCode.incompatible_return_value,
+                        detail=can_assign.display(),
+                    )
+        elif self.expected_return_value is not None:
             can_assign = self.expected_return_value.can_assign(value, self)
             if isinstance(can_assign, CanAssignError):
                 self._show_error_if_checking(
diff --git a/pyanalyze/signature.py b/pyanalyze/signature.py
@@ -59,6 +59,7 @@
 )
 from .typevar import resolve_bounds_map
 from .value import (
+    TypeVarLike,
     annotate_value,
     AnnotatedValue,
     AnySource,
@@ -522,10 +523,10 @@ class Signature:
     """Whether type checking can call the actual function to retrieve a precise return value."""
     evaluator: Optional[Evaluator] = None
     """Type evaluator for this function."""
-    typevars_of_params: Dict[str, List["TypeVar"]] = field(
+    typevars_of_params: Dict[str, List[TypeVarLike]] = field(
         init=False, default_factory=dict, repr=False, compare=False, hash=False
     )
-    all_typevars: Set["TypeVar"] = field(
+    all_typevars: Set[TypeVarLike] = field(
         init=False, default_factory=set, repr=False, compare=False, hash=False
     )
 
diff --git a/pyanalyze/test_generators.py b/pyanalyze/test_generators.py
@@ -0,0 +1,45 @@
+# static analysis: ignore
+from .implementation import assert_is_value
+from .test_name_check_visitor import TestNameCheckVisitorBase
+from .test_node_visitor import assert_passes
+from .value import AnySource, AnyValue, KnownValue, TypedValue
+
+
+class TestGenerator(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_generator_return(self):
+        from typing import Generator
+
+        def gen(cond) -> Generator[int, str, float]:
+            x = yield 1
+            assert_is_value(x, TypedValue(str))
+            yield "x"  # E: incompatible_yield
+            if cond:
+                return 3.0
+            else:
+                return "hello"  # E: incompatible_return_value
+
+        def capybara() -> Generator[int, int, int]:
+            x = yield from gen(True)  # E: incompatible_yield
+            assert_is_value(x, TypedValue(float))
+
+            return 3
+
+    @assert_passes()
+    def test_iterable_return(self):
+        from typing import Iterable
+
+        def gen(cond) -> Iterable[int]:
+            x = yield 1
+            assert_is_value(x, KnownValue(None))
+
+            yield "x"  # E: incompatible_yield
+
+            if cond:
+                return
+            else:
+                return 3  # E: incompatible_return_value
+
+        def caller() -> Iterable[int]:
+            x = yield from gen(True)
+            assert_is_value(x, AnyValue(AnySource.generic_argument))
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
@@ -2604,7 +2604,7 @@ def replace_known_sequence_value(value: Value) -> Value:
     return value
 
 
-def extract_typevars(value: Value) -> Iterable["TypeVar"]:
+def extract_typevars(value: Value) -> Iterable[TypeVarLike]:
     for val in value.walk_values():
         if isinstance(val, TypeVarValue):
             yield val.typevar

Original file line number	Diff line number	Diff line change
`@@ -59,6 +59,7 @@`
`59`	`59`	`)`
`60`	`60`	`from .typevar import resolve_bounds_map`
`61`	`61`	`from .value import (`
	`62`	`+ TypeVarLike,`
`62`	`63`	`annotate_value,`
`63`	`64`	`AnnotatedValue,`
`64`	`65`	`AnySource,`
`@@ -522,10 +523,10 @@ class Signature:`
`522`	`523`	`"""Whether type checking can call the actual function to retrieve a precise return value."""`
`523`	`524`	`evaluator: Optional[Evaluator] = None`
`524`	`525`	`"""Type evaluator for this function."""`
`525`		`- typevars_of_params: Dict[str, List["TypeVar"]] = field(`
	`526`	`+ typevars_of_params: Dict[str, List[TypeVarLike]] = field(`
`526`	`527`	`init=False, default_factory=dict, repr=False, compare=False, hash=False`
`527`	`528`	`)`
`528`		`- all_typevars: Set["TypeVar"] = field(`
	`529`	`+ all_typevars: Set[TypeVarLike] = field(`
`529`	`530`	`init=False, default_factory=set, repr=False, compare=False, hash=False`
`530`	`531`	`)`
`531`	`532`