Repository: quora/pyanalyze (python)

Issue:
with statements aren't typechecked
There is no `visit_With` or `visit_AsyncWith`.

Pull request description:
with and async with
Fixes #342

Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 3f23318..139d6fa 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Check `with` and `async with` statements (#344)
 - Improve type compatibility between generics and literals (#346)
 - Infer signatures for method wrapper objects (bound methods
   of builtin types) (#345)
diff --git a/pyanalyze/error_code.py b/pyanalyze/error_code.py
index 9a3b5e8..077e9d9 100644
--- a/pyanalyze/error_code.py
+++ b/pyanalyze/error_code.py
@@ -84,6 +84,7 @@ class ErrorCode(enum.Enum):
     missing_return = 67
     no_return_may_return = 68
     implicit_reexport = 69
+    invalid_context_manager = 70
 
 
 # Allow testing unannotated functions without too much fuss
@@ -191,6 +192,7 @@ ERROR_DESCRIPTION = {
     ErrorCode.missing_return: "Function may exit without returning a value",
     ErrorCode.no_return_may_return: "Function is annotated as NoReturn but may return",
     ErrorCode.implicit_reexport: "Use of implicitly re-exported name",
+    ErrorCode.invalid_context_manager: "Use of invalid object in with or async with",
 }
 
 
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index f4b002c..0cb598e 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -1412,11 +1412,15 @@ class NameCheckVisitor(
         except TypeError:
             return KnownValue(potential_function)
         if argspec is not None:
-            if info.async_kind != AsyncFunctionKind.async_proxy:
-                # don't attempt to infer the return value of async_proxy functions, since it will be
-                # set within the Future returned
-                # without this, we'll incorrectly infer the return value to be the Future instead of
-                # the Future's value
+            if (
+                info.async_kind != AsyncFunctionKind.async_proxy
+                and node.returns is None
+            ):
+                # Don't attempt to infer the return value of async_proxy functions, since it will be
+                # set within the Future returned. Without this, we'll incorrectly infer the return
+                # value to be the Future instead of the Future's value.
+                # Similarly, we don't infer the return value if there is an annotation, because
+                # we should be able to get it from __annotations__ instead.
                 self._argspec_to_retval[id(argspec)] = return_value
         else:
             self.log(logging.DEBUG, "No argspec", (potential_function, node))
@@ -3157,6 +3161,38 @@ class NameCheckVisitor(
             return AnyValue(AnySource.error)
         return result
 
+    def visit_With(self, node: ast.With) -> None:
+        for item in node.items:
+            self.visit_withitem(item)
+        self._generic_visit_list(node.body)
+
+    def visit_AsyncWith(self, node: ast.AsyncWith) -> None:
+        for item in node.items:
+            self.visit_withitem(item, is_async=True)
+        self._generic_visit_list(node.body)
+
+    def visit_withitem(self, node: ast.withitem, is_async: bool = False) -> None:
+        context = self.visit(node.context_expr)
+        if is_async:
+            protocol = "typing.AsyncContextManager"
+        else:
+            protocol = "typing.ContextManager"
+        val = GenericValue(protocol, [TypeVarValue(T)])
+        can_assign = val.can_assign(context, self)
+        if isinstance(can_assign, CanAssignError):
+            self._show_error_if_checking(
+                node.context_expr,
+                f"{context} is not a context manager",
+                detail=str(can_assign),
+                error_code=ErrorCode.invalid_context_manager,
+            )
+            assigned = AnyValue(AnySource.error)
+        else:
+            assigned = can_assign.get(T, AnyValue(AnySource.generic_argument))
+        if node.optional_vars is not None:
+            with qcore.override(self, "being_assigned", assigned):
+                self.visit(node.optional_vars)
+
     def visit_try_except(self, node: ast.Try) -> List[SubScope]:
         # reset yield checks between branches to avoid incorrect errors when we yield both in the
         # try and the except block
diff --git a/pyanalyze/signature.py b/pyanalyze/signature.py
index 427639f..28dd362 100644
--- a/pyanalyze/signature.py
+++ b/pyanalyze/signature.py
@@ -53,7 +53,7 @@ from .value import (
 
 import ast
 import asynq
-from collections import defaultdict, OrderedDict
+from collections import defaultdict
 import collections.abc
 from dataclasses import dataclass, field, replace
 import itertools
@@ -169,7 +169,7 @@ class CallContext:
     visitor: "NameCheckVisitor"
     """Using the visitor can allow various kinds of advanced logic
     in impl functions."""
-    bound_args: inspect.BoundArguments
+    composites: Dict[str, Composite]
     node: ast.AST
     """AST node corresponding to the function call. Useful for
     showing errors."""
@@ -193,7 +193,7 @@ class CallContext:
         return None
 
     def composite_for_arg(self, arg: str) -> Optional[Composite]:
-        composite = self.bound_args.arguments.get(arg)
+        composite = self.composites.get(arg)
         if isinstance(composite, Composite):
             return composite
         return None
@@ -408,8 +408,20 @@ class Signature:
             used_any = visitor.has_used_any_match()
         return tv_map, used_any
 
+    def _get_positional_parameter(self, index: int) -> Optional[SigParameter]:
+        for i, param in enumerate(self.signature.parameters.values()):
+            if param.kind in (
+                SigParameter.VAR_KEYWORD,
+                SigParameter.VAR_POSITIONAL,
+                SigParameter.KEYWORD_ONLY,
+            ):
+                return None
+            if i == index:
+                return param
+        return None
+
     def _apply_annotated_constraints(
-        self, raw_return: Union[Value, ImplReturn], bound_args: inspect.BoundArguments
+        self, raw_return: Union[Value, ImplReturn], composites: Dict[str, Composite]
     ) -> ImplReturn:
         if isinstance(raw_return, Value):
             ret = ImplReturn(raw_return)
@@ -420,12 +432,9 @@ class Signature:
             for guard in ret.return_value.get_metadata_of_type(
                 ParameterTypeGuardExtension
             ):
-                if guard.varname in bound_args.arguments:
-                    composite = bound_args.arguments[guard.varname]
-                    if (
-                        isinstance(composite, Composite)
-                        and composite.varname is not None
-                    ):
+                if guard.varname in composites:
+                    composite = composites[guard.varname]
+                    if composite.varname is not None:
                         constraint = Constraint(
                             composite.varname,
                             ConstraintType.is_value_object,
@@ -443,22 +452,21 @@ class Signature:
                     index = 1
                 else:
                     index = 0
-                composite = bound_args.args[index]
-                if isinstance(composite, Composite) and composite.varname is not None:
-                    constraint = Constraint(
-                        composite.varname,
-                        ConstraintType.is_value_object,
-                        True,
-                        guard.guarded_type,
-                    )
-                    constraints.append(constraint)
+                param = self._get_positional_parameter(index)
+                if param is not None:
+                    composite = composites[param.name]
+                    if composite.varname is not None:
+                        constraint = Constraint(
+                            composite.varname,
+                            ConstraintType.is_value_object,
+                            True,
+                            guard.guarded_type,
+                        )
+                        constraints.append(constraint)
             for guard in ret.return_value.get_metadata_of_type(HasAttrGuardExtension):
-                if guard.varname in bound_args.arguments:
-                    composite = bound_args.arguments[guard.varname]
-                    if (
-                        isinstance(composite, Composite)
-                        and composite.varname is not None
-                    ):
+                if guard.varname in composites:
+                    composite = composites[guard.varname]
+                    if composite.varname is not None:
                         constraint = Constraint(
                             composite.varname,
                             ConstraintType.add_annotation,
@@ -774,21 +782,13 @@ class Signature:
                 # You only get to do this once per call.
                 is_overload = False
 
-        inspect_bound_args = inspect.BoundArguments(
-            self.signature,
-            OrderedDict(
-                (param, composite) for param, (_, composite) in bound_args.items()
-            ),
-        )
+        composites = {param: composite for param, (_, composite) in bound_args.items()}
         # don't call the implementation function if we had an error, so that
         # the implementation function doesn't have to worry about basic
         # type checking
         if not had_error and self.impl is not None:
             ctx = CallContext(
-                vars=variables,
-                visitor=visitor,
-                bound_args=inspect_bound_args,
-                node=node,
+                vars=variables, visitor=visitor, composites=composites, node=node
             )
             return_value = self.impl(ctx)
 
@@ -806,7 +806,7 @@ class Signature:
         if return_value is EMPTY:
             ret = ImplReturn(AnyValue(AnySource.unannotated), is_error=had_error)
         else:
-            ret = self._apply_annotated_constraints(return_value, inspect_bound_args)
+            ret = self._apply_annotated_constraints(return_value, composites)
         return ret._replace(
             is_error=had_error,
             used_any_for_match=used_any,


Test patch:
diff --git a/pyanalyze/test_annotations.py b/pyanalyze/test_annotations.py
index a519c69..f12f9f0 100644
--- a/pyanalyze/test_annotations.py
+++ b/pyanalyze/test_annotations.py
@@ -140,7 +140,19 @@ class TestAnnotations(TestNameCheckVisitorBase):
             # Ideally should be ContextManager[int], but at least
             # it should not be Iterator[int], which is what pyanalyze
             # used to infer.
-            assert_is_value(capybara(), AnyValue(AnySource.inference))
+            assert_is_value(capybara(), AnyValue(AnySource.unannotated))
+
+            with capybara() as e:
+                assert_is_value(e, AnyValue(AnySource.generic_argument))
+
+            assert_is_value(post_capybara(), AnyValue(AnySource.unannotated))
+
+            with post_capybara() as e:
+                assert_is_value(e, AnyValue(AnySource.generic_argument))
+
+        @contextmanager
+        def post_capybara() -> Iterator[int]:
+            yield 3
 
     @assert_passes()
     def test_none_annotations(self):
diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
index d5a3cba..2a511e2 100644
--- a/pyanalyze/test_name_check_visitor.py
+++ b/pyanalyze/test_name_check_visitor.py
@@ -2579,6 +2579,66 @@ class TestWhile(TestNameCheckVisitorBase):
                     return 1
 
 
+class TestWith(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_with(self) -> None:
+        class BadCM1:
+            pass
+
+        class BadCM2:
+            def __enter__(self, extra_arg) -> int:
+                return 0
+
+            def __exit__(self, typ, value, tb):
+                pass
+
+        class GoodCM:
+            def __enter__(self) -> int:
+                return 0
+
+            def __exit__(self, typ, value, tb):
+                pass
+
+        def capybara():
+            with BadCM1() as e:  # E: invalid_context_manager
+                assert_is_value(e, AnyValue(AnySource.error))
+
+            with BadCM2() as e:  # E: invalid_context_manager
+                assert_is_value(e, AnyValue(AnySource.error))
+
+            with GoodCM() as e:
+                assert_is_value(e, TypedValue(int))
+
+    @assert_passes()
+    def test_async_with(self) -> None:
+        class BadCM1:
+            pass
+
+        class BadCM2:
+            async def __aenter__(self, extra_arg) -> int:
+                return 0
+
+            async def __aexit__(self, typ, value, tb):
+                pass
+
+        class GoodCM:
+            async def __aenter__(self) -> int:
+                return 0
+
+            async def __aexit__(self, typ, value, tb):
+                pass
+
+        async def capybara():
+            async with BadCM1() as e:  # E: invalid_context_manager
+                assert_is_value(e, AnyValue(AnySource.error))
+
+            async with BadCM2() as e:  # E: invalid_context_manager
+                assert_is_value(e, AnyValue(AnySource.error))
+
+            async with GoodCM() as e:
+                assert_is_value(e, TypedValue(int))
+
+
 class HasGetattr(object):
     def __getattr__(self, attr):
         return 42