Repository: quora/pyanalyze
Pull request #529 (2022-04-21 18:58:17)

Support for typechecking generators
- Support for typechecking generators
- fix self check



Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 71c90de..3e90de9 100644
--- 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
index 85aa56b..1dd6781 100644
--- a/pyanalyze/annotations.py
+++ b/pyanalyze/annotations.py
@@ -76,6 +76,7 @@ from .signature import (
 )
 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
index 5decbb4..3482536 100644
--- 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 @@ ERROR_DESCRIPTION = {
     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
index 701ad0f..d21a7ed 100644
--- a/pyanalyze/functions.py
+++ b/pyanalyze/functions.py
@@ -23,6 +23,7 @@ from .node_visitor import ErrorContext
 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,
@@ -30,6 +31,7 @@ from .value import (
     CanAssignContext,
     CanAssignError,
     GenericValue,
+    get_tv_map,
     KnownValue,
     SubclassValue,
     TypedValue,
@@ -43,6 +45,15 @@ FunctionDefNode = Union[ast.FunctionDef, ast.AsyncFunctionDef]
 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
index ea12bbc..f97a5c4 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -82,7 +82,12 @@ from .functions import (
     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 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         # 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 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             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 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
     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 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
     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 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             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 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         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 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         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 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             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
index de44b8e..64c2b57 100644
--- a/pyanalyze/signature.py
+++ b/pyanalyze/signature.py
@@ -59,6 +59,7 @@ from .type_evaluation import (
 )
 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/value.py b/pyanalyze/value.py
index 2c57836..26ff7da 100644
--- 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


Test patch:
diff --git a/pyanalyze/test_generators.py b/pyanalyze/test_generators.py
new file mode 100644
index 0000000..82ae2d5
--- /dev/null
+++ 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))