Add reverse CustomCheck (#383)

docs/changelog.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Allow `CustomCheck` to customize what values
+  a value can be assigned to (#383)
 - Fix incorrect inference of `self` argument on
   some nested methods (#382)
 - Fix compatibility between `Callable` and `Annotated`

pyanalyze/extensions.py

@@ -54,6 +54,12 @@ def func(arg: Annotated[str, LiteralOnly()]) -> None:
         func("x")  # ok
         func(str(some_call()))  # error
 
+    It is also possible to customize checks in the other direction
+    by overriding the ``can_be_assigned()`` method. For example, if
+    the above ``CustomCheck`` overrode the ``can_be_assigned`` method
+    instead, a value of type ``Annotated[str, LiteralOnly()]`` could
+    only be passed to functions that take a ``Literal`` parameter.
+
     A ``CustomCheck`` can also be generic over a ``TypeVar``. To implement support
     for ``TypeVar``, two more methods must be overridden:
 
@@ -64,7 +70,12 @@ def func(arg: Annotated[str, LiteralOnly()]) -> None:
 
     """
 
-    def can_assign(self, value: "Value", ctx: "CanAssignContext") -> "CanAssign":
+    def can_assign(self, __value: "Value", __ctx: "CanAssignContext") -> "CanAssign":
+        return {}
+
+    def can_be_assigned(
+        self, __value: "Value", __ctx: "CanAssignContext"
+    ) -> "CanAssign":
         return {}
 
     def walk_values(self) -> Iterable["Value"]:

pyanalyze/test_annotations.py

@@ -1150,6 +1150,33 @@ def caller(x: str) -> None:
             capybara("x" if x else "y")
             capybara("x" if x else x)  # E: incompatible_argument
 
+    @assert_passes()
+    def test_reverse_direction(self):
+        from pyanalyze.extensions import CustomCheck
+        from pyanalyze.value import (
+            CanAssignContext,
+            Value,
+            CanAssign,
+            flatten_values,
+            CanAssignError,
+        )
+        from typing import Any
+        from typing_extensions import Annotated
+
+        class DontAssignToAny(CustomCheck):
+            def can_be_assigned(self, value: Value, ctx: CanAssignContext) -> CanAssign:
+                for val in flatten_values(value, unwrap_annotated=True):
+                    if isinstance(val, AnyValue):
+                        return CanAssignError("Assignment to Any disallowed")
+                return {}
+
+        def want_any(x: Any) -> None:
+            pass
+
+        def capybara(arg: Annotated[str, DontAssignToAny()]) -> None:
+            want_any(arg)  # E: incompatible_argument
+            print(len(arg))
+
     @assert_passes()
     def test_no_any(self) -> None:
         from pyanalyze.extensions import NoAny

pyanalyze/value.py

@@ -95,9 +95,7 @@ def can_assign(self, other: "Value", ctx: "CanAssignContext") -> "CanAssign":
             if not tv_maps:
                 return CanAssignError(f"Cannot assign {other} to {self}")
             return unify_typevar_maps(tv_maps)
-        elif isinstance(other, AnnotatedValue):
-            return self.can_assign(other.value, ctx)
-        elif isinstance(other, TypeVarValue):
+        elif isinstance(other, (AnnotatedValue, TypeVarValue)):
             return other.can_be_assigned(self, ctx)
         elif (
             isinstance(other, UnboundMethodValue)
@@ -354,6 +352,8 @@ def __str__(self) -> str:
         return f"Any[{self.source.name}]"
 
     def can_assign(self, other: Value, ctx: CanAssignContext) -> CanAssign:
+        if isinstance(other, (AnnotatedValue, MultiValuedValue)):
+            return super().can_assign(other, ctx)
         return {}  # Always allowed
 
 
@@ -1793,6 +1793,18 @@ def can_assign(self, other: Value, ctx: CanAssignContext) -> CanAssign:
             tv_maps.append(custom_can_assign)
         return unify_typevar_maps(tv_maps)
 
+    def can_be_assigned(self, other: Value, ctx: CanAssignContext) -> CanAssign:
+        can_assign = other.can_assign(self.value, ctx)
+        if isinstance(can_assign, CanAssignError):
+            return can_assign
+        tv_maps = [can_assign]
+        for custom_check in self.get_metadata_of_type(CustomCheckExtension):
+            custom_can_assign = custom_check.custom_check.can_be_assigned(other, ctx)
+            if isinstance(custom_can_assign, CanAssignError):
+                return custom_can_assign
+            tv_maps.append(custom_can_assign)
+        return unify_typevar_maps(tv_maps)
+
     def walk_values(self) -> Iterable[Value]:
         yield self
         yield from self.value.walk_values()