Repository: quora/pyanalyze
Pull request #384 (2022-01-01 09:57:20)

More type evaluation improvements


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index d491e86..48766bb 100644
--- a/docs/changelog.md
+++ b/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`
@@ -12,7 +14,7 @@
 - Support the Python 3.10 `match` statement (#376)
 - Support the walrus (`:=`) operator (#375)
 - Initial support for proposed new "type evaluation"
-  mechanism (#374, #379)
+  mechanism (#374, #379, #384)
 - Create command-line options for each config option (#373)
 - Overhaul treatment of function definitions (#372)
   - Support positional-only arguments
diff --git a/docs/type_evaluation.md b/docs/type_evaluation.md
index f7fd0f3..2706805 100644
--- a/docs/type_evaluation.md
+++ b/docs/type_evaluation.md
@@ -1,7 +1,13 @@
 # Type evaluation
 
 Type evaluation is a mechanism for replacing complex
-overloads and version checks. For example, consider the
+overloads and version checks. It provides a restricted
+subset of Python that can be executed by type checkers
+to customize the behavior of a particular function.
+
+## Motivation
+
+Consider the
 definition of `round()` in typeshed:
 
     @overload
@@ -22,6 +28,33 @@ With type evaluation, this could instead be written as:
 
 This makes it easier to see at a glance what the difference is between various overloads.
 
+Other features of type evaluation, as proposed here, include
+customizable error messages, branching on the type of an
+argument, and branching on whether an argument was provided
+as a positional or keyword argument.
+
+Type evaluation functions can replace most complex overloads
+with simpler, more readable code. They solve a number of
+problems:
+
+- Error messages involving overloads are often hard to read.
+  Type evaluation functions enable the author of a function
+  to provide custom error messages that clearly point out
+  the issue in the user's code.
+- Complex overloads can be difficult to understand and write.
+  Type evaluation functions provide a more natural interface
+  that is closer to how such functions are written at
+  runtime.
+- The precise behavior of overloads is not specified and
+  varies across type checkers. The behavior of type
+  evaluation functions is more precisely specified.
+- Type evaluation functions provide ways to implement
+  several type system features that have been previously
+  requested, including:
+  - Marking a function, parameter, or parameter type as
+    deprecated.
+  - Accepting `Sequence[str]` but not `str`.
+
 ## Specification
 
 This section specifies how type evaluation works, without
@@ -49,7 +82,9 @@ do the following:
   evaluation function, as with a normal call.
 - Symbolically evaluate the body of the type evaluation
   until it reaches a `return` statement, which provides the type
-  that the call should return.
+  that the call should return. During this symbolic
+  evaluation, each argument is set to the value it has at
+  the call site that is being evaluated.
 - If execution reached a `return` statement, return the type
   provided by that statement. Otherwise, return the type
   set in the evaluation function's return annotation, or
@@ -61,6 +96,13 @@ a type annotation on a parameter or return type. However, no
 error should be provided if the return annotation is missing
 and all branches (including error branches) return a type.
 
+The default value of a parameter to an evaluation function
+may be either `...` or any value that is valid inside
+`Literal[...]`. If an argument with default `X` is not
+provided in a call, the type of the argument within the
+evaluation function is `Literal[X]`. If the default is
+`...`, the type is the parameter's annotation instead.
+
 Simple examples to demonstrate the semantics:
 
     @evaluated
@@ -85,6 +127,14 @@ Simple examples to demonstrate the semantics:
     x = always_errors(1)  # error
     reveal_type(x)  # str
 
+    @evaluated
+    def with_defaults(x: int = ..., y: int = 1) -> None:
+        reveal_type(x)
+        reveal_type(y)
+
+    with_defaults()  # x is "int", y is "Literal[1]"
+    with_defaults(1)  # x and y are both "Literal[1]"
+
 ### Supported features
 
 The body of a type evaluation uses a restricted subset of Python.
@@ -95,6 +145,10 @@ The only supported features are:
 - `pass` statements, which do nothing.
 - Calls to `show_error()`, which cause the type checker
   to emit an error. These are discussed further below.
+- Calls to `reveal_type(arg)`, where arg is one of the
+  arguments to the type evaluation function. These cause the
+  type checker to emit a message showing the current type
+  of `arg`. This is a debugging feature.
 
 Conditions in `if` statements may contain:
 
@@ -154,7 +208,7 @@ as a keyword argument.
 
 Parameters in Python can be provided in three ways,
 which we call _argument kinds_ for the purpose of this
-specificatioon:
+specification:
 
 - `POSITIONAL`: at the call site, either a single
   positional argument or a variadic one (`*args`)
@@ -335,9 +389,19 @@ taken, with the parameter type narrowed appropriately in each
 case. The return type of the function is the union of the two
 branches.
 
-For example, if the `ndigits` argument to
-`round()` is of type `int | None`, the inferred return value should
-be `_T | int`.
+For example:
+
+    @evaluated
+    def switch_types(arg: str | int):
+        if is_of_type(arg, str):
+            return int
+        else:
+            return str
+
+    reveal_type(switch_types(1))  # str
+    reveal_type(switch_types("x"))  # int
+    union: int | str
+    reveal_type(switch_types(union))  # int | str
 
 ### Type compatibility
 
@@ -523,6 +587,87 @@ This is how it could be implemented using `@evaluated`:
             show_error("start is a positional-only argument in Python <3.8", argument=start)
         return _T | _S
 
+## Possible extensions
+
+The following features may be useful, but are deferred
+for now for simplicity.
+
+### Error categories
+
+It may be useful to provide hints to the type checker
+about the severity of a `show_error()` call. For example,
+deprecation warnings could be marked so that the user can
+control whether to show them.
+
+One possibility is to add a keyword-only argument
+`category: str = ...` to `show_error()`. We would specify
+some standard categories that can be used in typeshed:
+
+- `deprecation` (for deprecated behavior)
+- `python_version` (for wrong Python version)
+- `platform` (for wrong sys.platform)
+- `warning` (for miscellaneous non-blocking issues)
+
+Type checkers could add support for additional categories
+as desired. Other type checkers would be expected to
+silently ignore unrecognized category strings.
+
+### Reusable error messages
+
+Because `show_error()` requires a string literal as the
+message, typeshed would contain a lot of hardcoded string
+messages about version changes.
+
+Some possible solutions include:
+
+- Allow the message to be a variable of `Literal` type
+  instead of a string literal. However, this would not
+  allow customizing an error message to include e.g.
+  the name of the argument or the Python version when
+  some behavior changed.
+- Allow the message to be a call to `.format()` on a
+  string literal or `Literal` variable, where all the
+  arguments are function arguments or literals:
+  `show_error(NEW_IN_VERSION.format(arg, "3.10"))`.
+- Allow the message to be a call to another evaluation
+  function that returns a string literal instead of a type.
+  This would allow even more complex logic for emitting the
+  error message.
+
+The last option could look like this:
+
+    @evaluated
+    def added_in_py_version(feature: str, version: str):
+        return f"{feature} was added in Python {version}"
+
+    def zip(strict: bool = False):
+        if is_provided(strict) and sys.version_info < (3, 10):
+            show_error(
+                added_in_py_version("strict", "3.10"), 
+                argument="strict"
+            )
+
+
+### Adding attributes
+
+A common pattern in type checker plugins is for the plugin
+to add some extra attribute to the object. For example,
+`@functools.total_ordering` inserts various dunder methods
+into the class it decorates.
+
+We could add an `add_attributes()` primitive that given
+a type and a dictionary of attributes, modifies the type
+to add these attributes.
+
+Usage could look like this:
+
+    @evaluated
+    def total_ordering(cls: Type[T]):
+        return add_attributes(
+            cls,
+            {"__eq__": Callable[[T, T], bool]}
+        )
+
 ## Status
 
 A partial implementation of this feature is available
@@ -550,6 +695,8 @@ Currently unsupported features include:
   an evaluation function for a runtime function,
   to replace some impls.
 - Type compatibility for evaluated functions.
+- Implementation of the desired behavior for
+  return annotations
 
 Areas that need more thought include:
 
@@ -557,23 +704,3 @@ Areas that need more thought include:
 - Interaction with overloads. It should be possible
   to register multiple evaluation functions for a
   function, treating them as overloads.
-- Consider adding support for `assert` and an
-  ergonomical way to produce a standardized error
-  if something is not supported in the current
-  version or platform.
-- Guidance on what the return annotation of an
-  evaluation function should be. Most likely,
-  it is treated as the default return type if
-  execution reaches the end of the evaluation
-  function. It can be omitted if the evaluation
-  function always return.
-- Add a `warn()` mechanism to warn on particular
-  invocations. This can be useful as a mechanism
-  to produce deprecation warnings.
-
-Motivations can include:
-
-- Less repetitive overload writing
-- Ability to customize error messages
-- Potential for additional features that work
-  across type checkers
diff --git a/pyanalyze/extensions.py b/pyanalyze/extensions.py
index 7815e53..7b6643d 100644
--- a/pyanalyze/extensions.py
+++ b/pyanalyze/extensions.py
@@ -54,6 +54,12 @@ class CustomCheck:
         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 @@ class CustomCheck:
 
     """
 
-    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"]:
diff --git a/pyanalyze/patma.py b/pyanalyze/patma.py
index cafbd9a..9a0ad86 100644
--- a/pyanalyze/patma.py
+++ b/pyanalyze/patma.py
@@ -24,9 +24,8 @@ from typing import (
     Union,
 )
 
-from pyanalyze.implementation import len_of_value
-from pyanalyze.signature import MappingValue
-
+from .implementation import len_of_value
+from .signature import MappingValue
 from .annotations import type_from_value
 from .extensions import CustomCheck
 from .error_code import ErrorCode
@@ -39,6 +38,7 @@ from .stacked_scopes import (
     ConstraintType,
     OrConstraint,
     constrain_value,
+    IsAssignablePredicate,
 )
 from .value import (
     AnnotatedValue,
@@ -61,6 +61,7 @@ from .value import (
     unannotate,
     unite_values,
     unpack_values,
+    is_overlapping,
     UNINITIALIZED_VALUE,
 )
 
@@ -137,25 +138,6 @@ MatchableSequence = AnnotatedValue(
 )
 
 
-@dataclass
-class IsAssignablePredicate:
-    pattern_value: Value
-    ctx: CanAssignContext
-    positive_only: bool
-
-    def __call__(self, value: Value, positive: bool) -> Optional[Value]:
-        compatible = is_overlapping(self.pattern_value, value, self.ctx)
-        if positive:
-            if not compatible:
-                return None
-            if value.is_assignable(self.pattern_value, self.ctx):
-                return self.pattern_value
-        elif not self.positive_only:
-            if compatible:
-                return None
-        return value
-
-
 @dataclass
 class LenPredicate:
     expected_length: int
@@ -450,10 +432,6 @@ class PatmaVisitor(ast.NodeVisitor):
             )
 
 
-def is_overlapping(left: Value, right: Value, ctx: CanAssignContext) -> bool:
-    return left.is_assignable(right, ctx) or right.is_assignable(left, ctx)
-
-
 def index_of(elts: Sequence[T], pred: Callable[[T], bool]) -> Optional[int]:
     for i, elt in enumerate(elts):
         if pred(elt):
diff --git a/pyanalyze/stacked_scopes.py b/pyanalyze/stacked_scopes.py
index 7aa456c..9ee78ee 100644
--- a/pyanalyze/stacked_scopes.py
+++ b/pyanalyze/stacked_scopes.py
@@ -67,6 +67,8 @@ from .value import (
     unite_and_simplify,
     unite_values,
     flatten_values,
+    CanAssignContext,
+    is_overlapping,
 )
 
 T = TypeVar("T")
@@ -514,6 +516,25 @@ class PredicateProvider(AbstractConstraint):
         return NULL_CONSTRAINT
 
 
+@dataclass
+class IsAssignablePredicate:
+    pattern_value: Value
+    ctx: CanAssignContext
+    positive_only: bool
+
+    def __call__(self, value: Value, positive: bool) -> Optional[Value]:
+        compatible = is_overlapping(self.pattern_value, value, self.ctx)
+        if positive:
+            if not compatible:
+                return None
+            if value.is_assignable(self.pattern_value, self.ctx):
+                return self.pattern_value
+        elif not self.positive_only:
+            if compatible:
+                return None
+        return value
+
+
 @dataclass(frozen=True)
 class AndConstraint(AbstractConstraint):
     """Represents the AND of two constraints."""
diff --git a/pyanalyze/type_evaluation.py b/pyanalyze/type_evaluation.py
index ad96cad..129681c 100644
--- a/pyanalyze/type_evaluation.py
+++ b/pyanalyze/type_evaluation.py
@@ -28,6 +28,13 @@ from typing import (
     Union,
 )
 
+from .stacked_scopes import (
+    Constraint,
+    ConstraintType,
+    VarnameWithOrigin,
+    constrain_value,
+    IsAssignablePredicate,
+)
 from .extensions import get_type_evaluation
 from .safe import all_of_type
 from .value import (
@@ -162,8 +169,6 @@ class IsOfTypeCondition(Condition):
         original_arg_type = self.original_arg_type
         remaining_type = self.remaining_type
         matched_type = subtract_unions(original_arg_type, remaining_type)
-        if negated:
-            matched_type, remaining_type = remaining_type, matched_type
 
         if matched_type is NO_RETURN_VALUE:
             text = negative_text
@@ -201,7 +206,7 @@ def subtract_unions(left: Value, right: Value) -> Value:
     return unite_values(*remaining)
 
 
-_Operator = Union[Type[ast.cmpop], Literal["matches", "does not match"]]
+_Operator = Union[Type[ast.cmpop], Literal["is of type", "is not of type"]]
 
 
 @dataclass
@@ -218,16 +223,16 @@ def _dummy_impl(left: object, right: object) -> object:
 _OP_TO_DATA: Dict[_Operator, _Comparator] = {
     ast.Is: _Comparator("is", ast.IsNot, operator.is_),
     ast.IsNot: _Comparator("is not", ast.Is, operator.is_not),
-    ast.Eq: _Comparator("==", ast.Eq, operator.eq),
-    ast.NotEq: _Comparator("!=", ast.NotEq, operator.ne),
+    ast.Eq: _Comparator("==", ast.NotEq, operator.eq),
+    ast.NotEq: _Comparator("!=", ast.Eq, operator.ne),
     ast.Gt: _Comparator(">", ast.LtE, operator.gt),
     ast.LtE: _Comparator("<=", ast.Gt, operator.le),
     ast.Lt: _Comparator("<", ast.GtE, operator.lt),
     ast.GtE: _Comparator(">=", ast.Lt, operator.ge),
     ast.In: _Comparator("in", ast.NotIn, lambda a, b: a in b),
     ast.NotIn: _Comparator("not in", ast.In, lambda a, b: a not in b),
-    "matches": _Comparator("matches", "matches", _dummy_impl),
-    "does not match": _Comparator("does not match", "matches", _dummy_impl),
+    "is of type": _Comparator("is of type", "is not of type", _dummy_impl),
+    "is not of type": _Comparator("is not of type", "is of type", _dummy_impl),
 }
 
 
@@ -401,7 +406,7 @@ class ConditionEvaluator(ast.NodeVisitor):
                         "Invalid keyword argument to is_of_type()", error_node
                     )
             return self.visit_is_of_type(
-                varname_node, typ, "matches", exclude_any=exclude_any
+                varname_node, typ, "is of type", exclude_any=exclude_any
             )
         else:
             return self.return_invalid(f"Invalid function {name}", node.func)
@@ -424,6 +429,17 @@ class ConditionEvaluator(ast.NodeVisitor):
         can_assign = can_assign_maybe_exclude_any(
             typ, val, self.ctx.can_assign_context, exclude_any
         )
+        condition = IsOfTypeCondition(
+            varname_node.id, op, val, NO_RETURN_VALUE, typ, exclude_any=exclude_any
+        )
+        constraint = Constraint(
+            VarnameWithOrigin(""),
+            ConstraintType.predicate,
+            True,
+            IsAssignablePredicate(
+                typ, self.ctx.can_assign_context, positive_only=False
+            ),
+        )
         if isinstance(can_assign, CanAssignError):
             pair = decompose_union(typ, val, self.ctx.can_assign_context, exclude_any)
             if pair is not None:
@@ -437,21 +453,14 @@ class ConditionEvaluator(ast.NodeVisitor):
                         typ,
                         exclude_any=exclude_any,
                     ),
-                    left_varmap={varname_node.id: typ},
+                    left_varmap={varname_node.id: constrain_value(val, constraint)},
                     right_varmap={varname_node.id: remaining},
                 )
-            return ConditionReturn(right_varmap={}, condition=NullCondition())
+            return ConditionReturn(right_varmap={}, condition=NotCondition(condition))
         else:
             return ConditionReturn(
-                left_varmap={varname_node.id: typ},
-                condition=IsOfTypeCondition(
-                    varname_node.id,
-                    op,
-                    val,
-                    NO_RETURN_VALUE,
-                    typ,
-                    exclude_any=exclude_any,
-                ),
+                left_varmap={varname_node.id: constrain_value(val, constraint)},
+                condition=condition,
             )
 
     def visit_UnaryOp(self, node: ast.UnaryOp) -> ConditionReturn:
@@ -525,10 +534,10 @@ class ConditionEvaluator(ast.NodeVisitor):
         with stack:
             for operand in node.values:
                 result = self.visit(operand)
+                active.append(result.condition)
                 if is_and:
                     if result.left_varmap is None:
                         # Condition returns False
-                        active.append(NotCondition(result.condition))
                         return ConditionReturn(
                             right_varmap=result.right_varmap,
                             condition=ConditionList(active),
@@ -539,10 +548,8 @@ class ConditionEvaluator(ast.NodeVisitor):
                         stack.enter_context(
                             self.ctx.narrow_variables(result.left_varmap)
                         )
-                        active.append(result.condition)
                     else:
                         # Condition matches partially
-                        active.append(result.condition)
                         narrowed_varmap.update(result.left_varmap)
                         stack.enter_context(
                             self.ctx.narrow_variables(result.left_varmap)
@@ -555,17 +562,14 @@ class ConditionEvaluator(ast.NodeVisitor):
                         stack.enter_context(
                             self.ctx.narrow_variables(result.right_varmap)
                         )
-                        active.append(NotCondition(result.condition))
                     elif result.right_varmap is None:
                         # Condition returns True
-                        active.append(result.condition)
                         return ConditionReturn(
                             left_varmap=result.left_varmap,
                             condition=ConditionList(active),
                         )
                     else:
                         # Condition partially matches
-                        active.append(result.condition)
                         narrowed_varmap.update(result.right_varmap)
                         stack.enter_context(
                             self.ctx.narrow_variables(result.right_varmap)
@@ -676,50 +680,69 @@ class EvaluateVisitor(ast.NodeVisitor):
         return self.ctx.evaluate_type(node.value)
 
     def visit_Expr(self, node: ast.Expr) -> EvalReturn:
-        if (
-            isinstance(node.value, ast.Call)
-            and isinstance(node.value.func, ast.Name)
-            and node.value.func.id == "show_error"
-        ):
-            call = node.value
-            if len(call.args) != 1:
-                self.add_invalid(
-                    "show_error() takes exactly one positional argument", call
-                )
-                return None
-            if not isinstance(call.args[0], ast.Str):
-                self.add_invalid(
-                    "show_error() message must be a string literal", call.args[0]
-                )
-                return None
-            message = call.args[0].s
-            argument = None
-            for keyword in call.keywords:
-                if keyword.arg == "argument":
-                    if not isinstance(keyword.value, ast.Name):
-                        self.add_invalid("argument must be a name", keyword.value)
-                        return None
-                    argument = keyword.value.id
-                    if argument not in self.ctx.variables:
-                        self.add_invalid(
-                            f"{argument} is not a valid argument", keyword.value
-                        )
-                        return None
-                else:
-                    # Before 3.9 keyword nodes don't have a lineno
-                    if sys.version_info >= (3, 9):
-                        error_node = keyword
-                    else:
-                        error_node = node
+        if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name):
+            name = node.value.func.id
+            if name == "show_error":
+                return self.visit_show_error(node.value)
+            elif name == "reveal_type":
+                return self.visit_reveal_type(node.value)
+        self.add_invalid("Invalid statement", node)
+        return None
+
+    def visit_reveal_type(self, call: ast.Call) -> EvalReturn:
+        if len(call.args) != 1 or call.keywords:
+            self.add_invalid(
+                "reveal_type() takes exactly one positional argument", call
+            )
+            return None
+        arg = call.args[0]
+        if not isinstance(arg, ast.Name):
+            self.add_invalid("reveal_type() argument must be a variable name", arg)
+            return None
+        try:
+            val = self.ctx.variables[arg.id]
+        except KeyError:
+            self.errors.append(InvalidEvaluation(f"Invalid variable {arg.id}", arg))
+            return None
+        message = (
+            f"Type of {arg.id} is {self.ctx.can_assign_context.display_value(val)}"
+        )
+        self.errors.append(UserRaisedError(message, [], argument=arg.id))
+        return None
+
+    def visit_show_error(self, call: ast.Call) -> EvalReturn:
+        if len(call.args) != 1:
+            self.add_invalid("show_error() takes exactly one positional argument", call)
+            return None
+        if not isinstance(call.args[0], ast.Str):
+            self.add_invalid(
+                "show_error() message must be a string literal", call.args[0]
+            )
+            return None
+        message = call.args[0].s
+        argument = None
+        for keyword in call.keywords:
+            if keyword.arg == "argument":
+                if not isinstance(keyword.value, ast.Name):
+                    self.add_invalid("argument must be a name", keyword.value)
+                    return None
+                argument = keyword.value.id
+                if argument not in self.ctx.variables:
                     self.add_invalid(
-                        "Invalid keyword argument to show_error()", error_node
+                        f"{argument} is not a valid argument", keyword.value
                     )
                     return None
-            self.errors.append(
-                UserRaisedError(message, list(self.active_conditions), argument)
-            )
-            return None
-        self.add_invalid("Invalid statement", node)
+            else:
+                # Before 3.9 keyword nodes don't have a lineno
+                if sys.version_info >= (3, 9):
+                    error_node = keyword
+                else:
+                    error_node = call
+                self.add_invalid("Invalid keyword argument to show_error()", error_node)
+                return None
+        self.errors.append(
+            UserRaisedError(message, list(self.active_conditions), argument)
+        )
         return None
 
     def visit_If(self, node: ast.If) -> EvalReturn:
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index 4991e3b..d598d44 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -95,9 +95,7 @@ class Value:
             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)
@@ -249,6 +247,10 @@ class CanAssignContext(Protocol):
         """Whether Any should be compatible only with itself."""
         return False
 
+    def display_value(self, value: Value) -> str:
+        """Provide a pretty, user-readable display of this value."""
+        return str(value)
+
 
 @dataclass(frozen=True)
 class CanAssignError:
@@ -354,6 +356,8 @@ class AnyValue(Value):
         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 +1797,18 @@ class AnnotatedValue(Value):
             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()
@@ -2312,3 +2328,7 @@ def can_assign_and_used_any(
         tv_map = param_typ.can_assign(var_value, ctx)
         used_any = ctx.has_used_any_match()
     return tv_map, used_any
+
+
+def is_overlapping(left: Value, right: Value, ctx: CanAssignContext) -> bool:
+    return left.is_assignable(right, ctx) or right.is_assignable(left, ctx)


Test patch:
diff --git a/pyanalyze/test_annotations.py b/pyanalyze/test_annotations.py
index 2e29db1..24df5e8 100644
--- a/pyanalyze/test_annotations.py
+++ b/pyanalyze/test_annotations.py
@@ -1150,6 +1150,33 @@ class TestCustomCheck(TestNameCheckVisitorBase):
             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
diff --git a/pyanalyze/test_type_evaluation.py b/pyanalyze/test_type_evaluation.py
index 33c0a3b..0763bb9 100644
--- a/pyanalyze/test_type_evaluation.py
+++ b/pyanalyze/test_type_evaluation.py
@@ -244,6 +244,58 @@ class TestTypeEvaluation(TestNameCheckVisitorBase):
         def capybara():
             assert_is_value(is_walrus_available(), KnownValue(expected))
 
+    @assert_passes()
+    def test_nested_ifs(self):
+        from pyanalyze.extensions import evaluated, is_of_type
+        from typing_extensions import Literal
+
+        @evaluated
+        def is_int(i: int):
+            if is_of_type(i, Literal[1, 2]):
+                if i == 1:
+                    return Literal[1]
+                elif i == 2:
+                    return Literal[2]
+            return Literal[3]
+
+        def capybara():
+            assert_is_value(is_int(1), KnownValue(1))
+
+    @assert_passes()
+    def test_not_equals(self):
+        from pyanalyze.extensions import evaluated
+
+        @evaluated
+        def want_one(x: int = 1, y: bool = True):
+            if x != 1:
+                show_error("want one", argument=x)
+            if y is not True:
+                show_error("want one", argument=y)
+            return None
+
+        def want_one(x: int = 1, y: bool = True) -> None:
+            pass
+
+        def capybara():
+            want_one(2)  # E: incompatible_call
+            want_one(y=False)  # E: incompatible_call
+
+    @assert_passes()
+    def test_reveal_type(self):
+        from pyanalyze.extensions import evaluated
+
+        @evaluated
+        def has_default(x: int = 1):
+            reveal_type(x)
+            return None
+
+        def has_default(x: int = 1) -> None:
+            pass
+
+        def capybara(i: int):
+            has_default()  # E: incompatible_call
+            has_default(i)  # E: incompatible_call
+
 
 class TestBoolOp(TestNameCheckVisitorBase):
     @assert_passes()
@@ -292,6 +344,26 @@ class TestBoolOp(TestNameCheckVisitorBase):
             assert_is_value(use_or(x_or_y), TypedValue(str))
             assert_is_value(use_or(x_or_z), TypedValue(str) | TypedValue(int))
 
+    @assert_passes()
+    def test_literal_or(self):
+        from pyanalyze.extensions import evaluated
+        from typing import Union
+
+        @evaluated
+        def is_one(i: int):
+            if i == 1 or i == -1:
+                show_error("bad argument", argument=i)
+                return int
+            return str
+
+        def is_one(i: int) -> Union[int, str]:
+            raise NotImplementedError
+
+        def capybara():
+            val = is_one(-1)  # E: incompatible_call
+            assert_is_value(val, TypedValue(int))
+            assert_is_value(is_one(2), TypedValue(str))
+
 
 class TestValidation(TestNameCheckVisitorBase):
     @assert_passes()
@@ -339,3 +411,87 @@ class TestValidation(TestNameCheckVisitorBase):
 
         def bad_evaluator(a: int) -> None:
             pass
+
+
+class TestExamples(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_open(self):
+        from pyanalyze.extensions import evaluated, is_of_type
+        from typing import Callable, Union, IO, BinaryIO, Any, Optional
+        from typing_extensions import Literal
+        from io import (
+            BufferedRandom,
+            BufferedReader,
+            BufferedWriter,
+            FileIO,
+            TextIOWrapper,
+        )
+
+        _OpenFile = Union[str, bytes, int]
+        _Opener = Callable[[str, int], int]
+
+        # These are simplified
+        OpenTextModeUpdating = Literal["r+", "w+", "a+", "x+"]
+        OpenTextModeWriting = Literal["w", "wt", "tw", "a", "at", "ta", "x", "xt", "tx"]
+        OpenTextModeReading = Literal["r", "rt", "tr"]
+        OpenTextMode = Union[
+            OpenTextModeUpdating, OpenTextModeWriting, OpenTextModeReading
+        ]
+        OpenBinaryModeUpdating = Literal["rb+", "wb+", "ab+", "xb+"]
+        OpenBinaryModeWriting = Literal["wb", "bw", "ab", "ba", "xb", "bx"]
+        OpenBinaryModeReading = Literal["rb", "br"]
+        OpenBinaryMode = Union[
+            OpenBinaryModeUpdating, OpenBinaryModeReading, OpenBinaryModeWriting
+        ]
+
+        @evaluated
+        def open2(
+            file: _OpenFile,
+            mode: str = "r",
+            buffering: int = -1,
+            encoding: Optional[str] = None,
+            errors: Optional[str] = None,
+            newline: Optional[str] = None,
+            closefd: bool = False,
+            opener: Optional[_Opener] = None,
+        ) -> IO[Any]:
+            if is_of_type(mode, OpenTextMode):
+                return TextIOWrapper
+            elif is_of_type(mode, OpenBinaryMode):
+                if encoding is not None:
+                    show_error(
+                        "'encoding' argument may not be provided in binary moode",
+                        argument=encoding,
+                    )
+                if errors is not None:
+                    show_error(
+                        "'errors' argument may not be provided in binary moode",
+                        argument=errors,
+                    )
+                if newline is not None:
+                    show_error(
+                        "'newline' argument may not be provided in binary moode",
+                        argument=newline,
+                    )
+                if buffering == 0:
+                    return FileIO
+                elif buffering == -1 or buffering == 1:
+                    if is_of_type(mode, OpenBinaryModeUpdating):
+                        return BufferedRandom
+                    elif is_of_type(mode, OpenBinaryModeWriting):
+                        return BufferedWriter
+                    elif is_of_type(mode, OpenBinaryModeReading):
+                        return BufferedReader
+
+                # Buffering cannot be determined: fall back to BinaryIO
+                return BinaryIO
+            # Fallback if mode is not specified
+            return IO[Any]
+
+        def capybara():
+            assert_is_value(open2("x", "r"), TypedValue(TextIOWrapper))
+            open2("x", "rb", encoding="utf-8")  # E: incompatible_call
+            assert_is_value(open2("x", "rb", buffering=0), TypedValue(FileIO))
+            assert_is_value(open2("x", "rb+"), TypedValue(BufferedRandom))
+            assert_is_value(open2("x", "rb"), TypedValue(BufferedReader))
+            assert_is_value(open2("x", "rb", buffering=1), TypedValue(BufferedReader))