Repository: quora/pyanalyze
Pull request #374 (2021-12-28 11:58:57)

New "type evaluation" mechanism


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index c12be5d..7a3a321 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Initial support for proposed new "type evaluation"
+  mechanism (#374)
 - Create command-line options for each config option (#373)
 - Overhaul treatment of function definitions (#372)
   - Support positional-only arguments
diff --git a/docs/index.rst b/docs/index.rst
index 5173902..029743a 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -14,6 +14,7 @@ pyanalyze: A semi-static typechecker
    configuration
    typesystem
    design
+   type_evaluation
    glossary
    reference/annotations
    reference/ast_annotator
diff --git a/docs/type_evaluation.md b/docs/type_evaluation.md
new file mode 100644
index 0000000..dbf37c2
--- /dev/null
+++ b/docs/type_evaluation.md
@@ -0,0 +1,139 @@
+# Type evaluation
+
+Type evaluation is a mechanism for replacing complex
+overloads and version checks. For example, consider the
+definition of `round()` in typeshed:
+
+    @overload
+    def round(number: SupportsRound[Any]) -> int: ...
+    @overload
+    def round(number: SupportsRound[Any], ndigits: None) -> int: ...
+    @overload
+    def round(number: SupportsRound[_T], ndigits: SupportsIndex) -> _T: ...
+
+With type evaluation, this could instead be written as:
+
+    @evaluated
+    def round(number: SupportsRound[_T], ndigits: SupportsIndex | None = ...):
+        if not is_set(ndigits) or ndigits is None:
+            return int
+        else:
+            return _T
+
+This makes it easier to see at a glance what the difference is between various overloads.
+
+## Specification
+
+Type-evaluated functions may be declared both at runtime and in stub files, similar to the existing `@overload` mechanism. At runtime, the evaluated function must immediately precede
+the function implementation:
+
+        @evaluated
+        def round(number: SupportsRound[_T], ndigits: SupportsIndex | None = ...):
+            if ...
+
+        def round(number, ndigits=None):
+            return number.__round__(ndigits)
+
+In stubs, the implementation is omitted.
+
+When a type checker encounters a call to a function for which a type evaluation has been provided, it should symbolically evaluate the body of the type evaluation until it reaches a `return` or `raise` statement. A `raise` statement indicates that the type checker should produce an error, a `return` statement provides the type that the call should return.
+
+### Supported features
+
+The body of a type evaluation uses a restricted subset of Python.
+The only supported features are:
+
+- `if` statements and `else` blocks. These can only contain conditions of the form specified below.
+- `return` statements with return values that are interpretable as type annotations. This indicates the type that the function returns in a particular condition.
+- `raise` statements of the form `raise Exception(message)`, where `message` is a string literal. When the code takes a branch that ends in a `raise` statement, the type checker should emit an error with the provided message.
+
+Conditions in `if` statements may contain:
+- Calls to the special `is_set()` function, which takes as its argument a variable. This functions returns True if the variable was provided as an argument in the call. For example, given a function `def f(arg=None): ...`, `is_set(arg)` would return True if the function is invoked as `f(None)`, but False if it is invoked as `f()`. A dummy runtime implementation of `is_set()` is provided.
+- Calls to `isinstance()`, but with an extended meaning. `isinstance(arg, type)` will return True if and only if the type checker would accept an assignment `_: type = arg`. For example, `isinstance(arg, Literal["a", "b"])` is valid and returns True if the type of the argument is (for example) `Literal["a"]`.
+- Expressions of the form `arg is (not) <constant>`, where `<constant>` may be True, False, or None.
+- Expressions of the form `arg == <constant>` or `arg != <constant>`, where `<constant>` is any value valid inside `Literal[]` (a bool, int, string, or enum member).
+- Version and platform checks that are otherwise valid in stubs, as specified in PEP 484.
+- Multiple conditions combined with `and` or `or`.
+- A negation of another condition with `not`.
+
+### Interaction with Any
+
+What should `round()` return if the type of the `ndigits`
+argument is `Any`? Existing type checkers do not agree:
+pyright picks the first overload that matches and returns
+`int`, since `Any` is compatible with `None`; mypy and pyanalyze
+see that multiple overloads might match and return `Any`. There
+are good reasons for both choices<!-- insert link to Eric's explanation-->,
+and we allow the same behavior for type evaluations.
+
+Type checkers should pick one of the following two behaviors and 
+document their choice:
+1. All checks (`isinstance`, `is`, `==`) against variables typed 
+   as `Any` in the body of type evaluation succeed. 
+   `round(..., Any)` returns `int`. Note that
+   this means that switching the `if` and `else` blocks may change
+   visible behavior.
+2. Conditions on variables typed as `Any` take both branches of the
+   conditional. If the two branches return different types, `Any`
+   is returned instead. `round(..., Any)` returns `Any`.
+
+### Interaction with unions
+
+Type checkers should apply normal type narrowing rules to arguments
+that are of Union types. For example, if the `ndigits` argument to
+`round()` is of type `int | None`, the inferred return value should
+be `_T | int`.
+
+## Status
+
+A partial implementation of this feature is available
+in pyanalyze:
+
+    from pyanalyze.extensions import evaluated, is_set
+
+    @evaluated
+    def simple_evaluated(x: int, y: str = ""):
+        if is_set(y):
+            return int
+        else:
+            return str
+
+    def simple_evaluated(*args: object) -> Union[int, str]:
+        if len(args) >= 2:
+            return 1
+        else:
+            return "x"
+
+Currently unsupported features include:
+- Comparison against enum members
+- Version and platform checks
+- Use of `and` and `or`
+- Usage in stubs
+- pyanalyze should provide a way to register
+  an evaluation function for a runtime function,
+  to replace some impls.
+
+Areas that need more thought include:
+- Interaction with typevars
+- 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/__init__.py b/pyanalyze/__init__.py
index 802f6e2..a70c77d 100644
--- a/pyanalyze/__init__.py
+++ b/pyanalyze/__init__.py
@@ -45,6 +45,8 @@ used(dump_value)
 used(extensions.LiteralOnly)
 used(extensions.NoAny)
 used(extensions.overload)
+used(extensions.evaluated)
+used(extensions.is_set)
 used(value.UNRESOLVED_VALUE)  # keeping it around for now just in case
 used(reexport)
 used(checker)
diff --git a/pyanalyze/annotations.py b/pyanalyze/annotations.py
index 4325e4e..e064296 100644
--- a/pyanalyze/annotations.py
+++ b/pyanalyze/annotations.py
@@ -61,11 +61,13 @@ from .extensions import (
 from .find_unused import used
 from .signature import SigParameter, Signature, ParameterKind
 from .safe import is_typing_name, is_instance_of_typing_name
+from . import type_evaluation
 from .value import (
     AnnotatedValue,
     AnySource,
     AnyValue,
     CallableValue,
+    CanAssignContext,
     CustomCheckExtension,
     Extension,
     HasAttrGuardExtension,
@@ -144,6 +146,21 @@ class Context:
         return self.handle_undefined_name(name)
 
 
+@dataclass
+class TypeEvaluationContext(Context, type_evaluation.Context):
+    variables: type_evaluation.VarMap
+    set_variables: Container[str]
+    can_assign_context: CanAssignContext
+    globals: Mapping[str, object]
+
+    def evaluate_type(self, node: ast.AST) -> Value:
+        return type_from_ast(node, ctx=self)
+
+    def get_name(self, node: ast.Name) -> Value:
+        """Return the :class:`Value <pyanalyze.value.Value>` corresponding to a name."""
+        return self.get_name_from_globals(node.id, self.globals)
+
+
 @used  # part of an API
 def type_from_ast(
     ast_node: ast.AST,
diff --git a/pyanalyze/arg_spec.py b/pyanalyze/arg_spec.py
index e2d97a6..0eecdc9 100644
--- a/pyanalyze/arg_spec.py
+++ b/pyanalyze/arg_spec.py
@@ -32,6 +32,7 @@ from .signature import (
     Signature,
     ParameterKind,
 )
+from .type_evaluation import get_evaluator
 from .typeshed import TypeshedFinder
 from .value import (
     AnySource,
@@ -54,7 +55,7 @@ import asyncio
 import asynq
 from collections.abc import Awaitable
 import contextlib
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 import qcore
 import inspect
 import sys
@@ -507,6 +508,13 @@ class ArgSpecCache:
                     ]
                     if all_of_type(sigs, Signature):
                         return OverloadedSignature(sigs)
+                evaluator = get_evaluator(obj)
+                if evaluator is not None:
+                    sig = self._cached_get_argspec(
+                        evaluator.func, impl, is_asynq, in_overload_resolution=True
+                    )
+                    if isinstance(sig, Signature):
+                        return replace(sig, evaluator=evaluator)
 
         if isinstance(obj, tuple) or hasattr(obj, "__getattr__"):
             return None  # lost cause
diff --git a/pyanalyze/extensions.py b/pyanalyze/extensions.py
index fad52dd..5c945e6 100644
--- a/pyanalyze/extensions.py
+++ b/pyanalyze/extensions.py
@@ -17,6 +17,7 @@ from typing import (
     Container,
     Dict,
     Iterable,
+    Optional,
     Tuple,
     List,
     Union,
@@ -371,6 +372,7 @@ def reveal_type(value: object) -> None:
 
 
 _overloads: Dict[str, List[Callable[..., Any]]] = defaultdict(list)
+_type_evaluations: Dict[str, Optional[Callable[..., Any]]] = {}
 
 
 def get_overloads(fully_qualified_name: str) -> List[Callable[..., Any]]:
@@ -378,6 +380,11 @@ def get_overloads(fully_qualified_name: str) -> List[Callable[..., Any]]:
     return _overloads[fully_qualified_name]
 
 
+def get_type_evaluation(fully_qualified_name: str) -> Optional[Callable[..., Any]]:
+    """Return the type evaluation function for this fully qualified name, or None."""
+    return _type_evaluations.get(fully_qualified_name)
+
+
 if TYPE_CHECKING:
     from typing import overload as overload
 
@@ -394,3 +401,21 @@ else:
         key = f"{func.__module__}.{func.__qualname__}"
         _overloads[key].append(func)
         return real_overload(func)
+
+
+def evaluated(func: Callable[..., Any]) -> Callable[..., Any]:
+    """Marks a type evaluation function."""
+    key = f"{func.__module__}.{func.__qualname__}"
+    assert key not in _type_evaluations, f"multiple evaluations for {key}"
+    _type_evaluations[key] = func
+    func.__is_type_evaluation__ = True
+    return func
+
+
+def is_set(argument: object) -> bool:
+    """Helper function for type evaluators.
+
+    May not be called at runtime.
+
+    """
+    raise NotImplementedError("Should only be called in type evaluation functions")
diff --git a/pyanalyze/signature.py b/pyanalyze/signature.py
index 33b540f..6a0b0d5 100644
--- a/pyanalyze/signature.py
+++ b/pyanalyze/signature.py
@@ -6,6 +6,8 @@ calls.
 
 """
 
+from . import annotations
+from .type_evaluation import Evaluator, decompose_union, evaluate
 from .error_code import ErrorCode
 from .safe import all_of_type
 from .stacked_scopes import (
@@ -31,7 +33,6 @@ from .value import (
     HasAttrGuardExtension,
     KVPair,
     KnownValue,
-    MultiValuedValue,
     NoReturnConstraintExtension,
     NoReturnGuardExtension,
     ParameterTypeGuardExtension,
@@ -46,6 +47,7 @@ from .value import (
     CanAssign,
     CanAssignError,
     annotate_value,
+    can_assign_and_used_any,
     concrete_values_from_iterable,
     extract_typevars,
     flatten_values,
@@ -93,6 +95,7 @@ EMPTY = inspect.Parameter.empty
 UNANNOTATED = AnyValue(AnySource.unannotated)
 ARGS = qcore.MarkerObject("*args")
 KWARGS = qcore.MarkerObject("**kwargs")
+DEFAULT = qcore.MarkerObject("default")
 
 
 @dataclass
@@ -109,8 +112,8 @@ Argument = Tuple[Composite, Union[None, str, Literal[ARGS], Literal[KWARGS]]]
 
 # Arguments bound to a call. The first member of the tuple is a reference back
 # to the ActualArguments: the value is a str for a kwarg, int for a positional,
-# None for something we cannot infer.
-BoundArgs = Dict[str, Tuple[Union[None, str, int], Composite]]
+# DEFAULT for something filled from the defaults, or None for something we cannot infer.
+BoundArgs = Dict[str, Tuple[Union[None, Literal[DEFAULT], str, int], Composite]]
 
 
 @dataclass
@@ -350,6 +353,8 @@ class Signature:
     a callable is compatible with any other callable with a compatible return type."""
     allow_call: bool = False
     """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(
         init=False, default_factory=dict, repr=False, compare=False, hash=False
     )
@@ -395,31 +400,14 @@ class Signature:
                 param_typ = param.annotation.substitute_typevars(typevar_map)
             else:
                 param_typ = param.annotation
-            tv_map, used_any = self._can_assign_and_used_any(
+            tv_map, used_any = can_assign_and_used_any(
                 param_typ, composite.value, visitor
             )
             if isinstance(tv_map, CanAssignError):
-                if is_overload and isinstance(composite.value, MultiValuedValue):
-                    tv_maps = []
-                    remaining_values = []
-                    union_used_any = False
-                    for val in composite.value.vals:
-                        can_assign, subval_used_any = self._can_assign_and_used_any(
-                            param_typ, val, visitor
-                        )
-                        if isinstance(can_assign, CanAssignError):
-                            remaining_values.append(val)
-                        else:
-                            if subval_used_any:
-                                union_used_any = True
-                            tv_maps.append(can_assign)
-                    if tv_maps:
-                        tv_map = unify_typevar_maps(tv_maps)
-                        assert remaining_values, (
-                            f"all union members matched between {param_typ} and"
-                            f" {composite.value}"
-                        )
-                        return tv_map, union_used_any, unite_values(*remaining_values)
+                if is_overload:
+                    triple = decompose_union(param_typ, composite.value, visitor)
+                    if triple is not None:
+                        return triple
                 visitor.show_error(
                     composite.node if composite.node is not None else node,
                     f"Incompatible argument type for {param.name}: expected {param_typ}"
@@ -431,14 +419,6 @@ class Signature:
             return tv_map, used_any, None
         return {}, False, None
 
-    def _can_assign_and_used_any(
-        self, param_typ: Value, var_value: Value, visitor: "NameCheckVisitor"
-    ) -> Tuple[CanAssign, bool]:
-        with visitor.reset_any_used():
-            tv_map = param_typ.can_assign(var_value, visitor)
-            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.parameters.values()):
             if param.kind in (
@@ -573,7 +553,7 @@ class Signature:
                     )
                     return None
                 else:
-                    bound_args[param.name] = None, Composite(param.default)
+                    bound_args[param.name] = DEFAULT, Composite(param.default)
             elif param.kind is ParameterKind.POSITIONAL_OR_KEYWORD:
                 if positional_index < len(actual_args.positionals):
                     bound_args[param.name] = (
@@ -602,7 +582,7 @@ class Signature:
                     # It may also come from **kwargs
                     if actual_args.star_kwargs is not None:
                         value = unite_values(
-                            actual_args.star_kwargs, actual_args.star_kwargs
+                            actual_args.star_args, actual_args.star_kwargs
                         )
                         star_kwargs_consumed = True
                     else:
@@ -629,7 +609,7 @@ class Signature:
                     )
                     return None
                 else:
-                    bound_args[param.name] = None, Composite(param.default)
+                    bound_args[param.name] = DEFAULT, Composite(param.default)
             elif param.kind is ParameterKind.KEYWORD_ONLY:
                 if param.name in actual_args.keywords:
                     definitely_provided, composite = actual_args.keywords[param.name]
@@ -653,7 +633,7 @@ class Signature:
                     )
                     return None
                 else:
-                    bound_args[param.name] = None, Composite(param.default)
+                    bound_args[param.name] = DEFAULT, Composite(param.default)
             elif param.kind is ParameterKind.VAR_POSITIONAL:
                 star_args_consumed = True
                 positionals = []
@@ -843,11 +823,34 @@ class Signature:
         # 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, composites=composites, node=node
-            )
-            return_value = self.impl(ctx)
+        if not had_error:
+            if self.impl is not None:
+                ctx = CallContext(
+                    vars=variables, visitor=visitor, composites=composites, node=node
+                )
+                return_value = self.impl(ctx)
+            elif self.evaluator is not None:
+                varmap = {
+                    param: composite.value
+                    for param, (_, composite) in bound_args.items()
+                }
+                set_variables = {
+                    param
+                    for param, (position, _) in bound_args.items()
+                    if position is not DEFAULT
+                }
+                ctx = annotations.TypeEvaluationContext(
+                    varmap, set_variables, visitor, self.evaluator.globals
+                )
+                return_value = evaluate(self.evaluator, ctx)
+                if isinstance(return_value, CanAssignError):
+                    self.show_call_error(
+                        "Error in type evaluator",
+                        node,
+                        visitor,
+                        detail=str(return_value),
+                    )
+                    return_value = AnyValue(AnySource.error)
 
         if self.allow_call:
             runtime_return = self._maybe_perform_call(preprocessed, visitor, node)
@@ -1215,6 +1218,7 @@ class Signature:
             has_return_annotation=self.has_return_annotation,
             is_ellipsis_args=is_ellipsis_args,
             allow_call=self.allow_call,
+            evaluator=self.evaluator,
         )
 
     def walk_values(self) -> Iterable[Value]:
@@ -1237,6 +1241,7 @@ class Signature:
             is_ellipsis_args=self.is_ellipsis_args,
             is_asynq=False,
             allow_call=self.allow_call,
+            evaluator=self.evaluator,
         )
 
     @classmethod
@@ -1251,6 +1256,7 @@ class Signature:
         is_ellipsis_args: bool = False,
         is_asynq: bool = False,
         allow_call: bool = False,
+        evaluator: Optional[Evaluator] = None,
     ) -> "Signature":
         """Create a :class:`Signature` object.
 
@@ -1271,12 +1277,17 @@ class Signature:
             is_ellipsis_args=is_ellipsis_args,
             is_asynq=is_asynq,
             allow_call=allow_call,
+            evaluator=evaluator,
         )
 
     def __str__(self) -> str:
         param_str = ", ".join(self._render_parameters())
         asynq_str = "@asynq " if self.is_asynq else ""
         rendered = f"{asynq_str}({param_str}) -> {self.return_value}"
+        if self.impl:
+            rendered += " (with impl)"
+        if self.evaluator:
+            rendered += " (with evaluator)"
         return rendered
 
     def _render_parameters(self) -> Iterable[str]:
diff --git a/pyanalyze/type_evaluation.py b/pyanalyze/type_evaluation.py
new file mode 100644
index 0000000..fd014c5
--- /dev/null
+++ b/pyanalyze/type_evaluation.py
@@ -0,0 +1,366 @@
+"""
+
+Implementation of type evaluation.
+
+"""
+
+import ast
+from contextlib import contextmanager
+from dataclasses import dataclass
+import inspect
+import textwrap
+from typing import (
+    Any,
+    Callable,
+    Container,
+    Iterator,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
+
+from .extensions import get_type_evaluation
+from .value import (
+    AnySource,
+    AnyValue,
+    CanAssignContext,
+    CanAssignError,
+    KnownValue,
+    MultiValuedValue,
+    Value,
+    can_assign_and_used_any,
+    unannotate,
+    unite_values,
+    unify_typevar_maps,
+    TypeVarMap,
+)
+
+VarMap = Mapping[str, Value]
+
+
+class InvalidEvaluation(Exception):
+    pass
+
+
+class Context:
+    variables: VarMap
+    set_variables: Container[str]
+    can_assign_context: CanAssignContext
+
+    def evaluate_type(self, __node: ast.AST) -> Value:
+        raise NotImplementedError
+
+    @contextmanager
+    def narrow_variables(self, varmap: Optional[VarMap]) -> Iterator[None]:
+        if varmap is None:
+            yield
+            return
+        old_varmap = self.variables
+        new_varmap = {**old_varmap, **varmap}
+        try:
+            self.variables = new_varmap
+            yield
+        finally:
+            self.variables = old_varmap
+
+
+@dataclass
+class Evaluator:
+    func: Callable[..., Any]
+    node: ast.FunctionDef
+    globals: Mapping[str, object]
+
+
+@dataclass
+class UnionCombinedReturn:
+    left: "EvalReturn"
+    right: "EvalReturn"
+
+
+@dataclass
+class AnyCombinedReturn:
+    left: "EvalReturn"
+    right: "EvalReturn"
+
+
+EvalReturn = Union[None, Value, CanAssignError, UnionCombinedReturn, AnyCombinedReturn]
+
+
+def may_be_none(ret: EvalReturn) -> bool:
+    if ret is None:
+        return True
+    elif isinstance(ret, (Value, CanAssignError)):
+        return False
+    else:
+        return may_be_none(ret.left) or may_be_none(ret.right)
+
+
+@dataclass
+class ConditionReturn:
+    # These are None if there is no match, and a (possibly empty)
+    # map of new variable values if there is a match.
+    left_varmap: Optional[VarMap] = None
+    right_varmap: Optional[VarMap] = None
+    is_any_match: bool = False
+
+    def reverse(self) -> "ConditionReturn":
+        return ConditionReturn(
+            left_varmap=self.right_varmap,
+            right_varmap=self.left_varmap,
+            is_any_match=self.is_any_match,
+        )
+
+
+@dataclass
+class ConditionEvaluator(ast.NodeVisitor):
+    ctx: Context
+
+    def visit_Call(self, node: ast.Call) -> ConditionReturn:
+        if not isinstance(node.func, ast.Name):
+            raise InvalidEvaluation("Unexpected call")
+        name = node.func.id
+        if name == "is_set":
+            if node.keywords or len(node.args) != 1:
+                raise InvalidEvaluation("is_set() takes a single argument")
+            if not isinstance(node.args[0], ast.Name):
+                raise InvalidEvaluation("Argument to is_set() must be a variable")
+            variable = node.args[0].id
+            match = variable in self.ctx.set_variables
+            if match:
+                return ConditionReturn(left_varmap={})
+            else:
+                return ConditionReturn(right_varmap={})
+        elif name == "isinstance":
+            if node.keywords or len(node.args) != 2:
+                raise InvalidEvaluation("isinstance() takes two positional arguments")
+            varname_node = node.args[0]
+            typ = self.ctx.evaluate_type(node.args[1])
+            return self.visit_isinstance(varname_node, typ)
+        else:
+            raise InvalidEvaluation(f"Invalid function {name}")
+
+    def visit_isinstance(self, varname_node: ast.AST, typ: Value) -> ConditionReturn:
+        if not isinstance(varname_node, ast.Name):
+            raise InvalidEvaluation("First argument to isinstance() must be a name")
+        val = self.get_name(varname_node)
+        can_assign, used_any = can_assign_and_used_any(
+            typ, val, self.ctx.can_assign_context
+        )
+        if isinstance(can_assign, CanAssignError):
+            triple = decompose_union(typ, val, self.ctx.can_assign_context)
+            if triple is not None:
+                _, used_any, remaining = triple
+                return ConditionReturn(
+                    left_varmap={varname_node.id: typ},
+                    right_varmap={varname_node.id: remaining},
+                    is_any_match=used_any,
+                )
+            return ConditionReturn(right_varmap={})
+        else:
+            return ConditionReturn(
+                left_varmap={varname_node.id: typ}, is_any_match=used_any
+            )
+
+    def visit_UnaryOp(self, node: ast.UnaryOp) -> ConditionReturn:
+        if isinstance(node.op, ast.Not):
+            ret = self.visit(node.operand)
+            return ret.reverse()
+        else:
+            raise InvalidEvaluation("Unsupported unary operation")
+
+    def visit_Compare(self, node: ast.Compare) -> ConditionReturn:
+        if len(node.ops) != 1:
+            raise InvalidEvaluation("Chained comparison is unsupported")
+        op = node.ops[0]
+        right = node.comparators[0]
+        if isinstance(op, (ast.Is, ast.IsNot)):
+            if not isinstance(right, ast.NameConstant):
+                raise InvalidEvaluation(
+                    "is/is not are only supported with True, False, and None"
+                )
+            ret = self.visit_isinstance(node.left, KnownValue(right.value))
+            if isinstance(op, ast.IsNot):
+                return ret.reverse()
+            return ret
+        elif isinstance(op, (ast.Eq, ast.NotEq)):
+            operand = self.evaluate_literal(right)
+            ret = self.visit_isinstance(node.left, operand)
+            if isinstance(op, ast.NotEq):
+                return ret.reverse()
+            return ret
+        else:
+            raise InvalidEvaluation("Unsupported comparison operator")
+
+    def evaluate_literal(self, node: ast.expr) -> KnownValue:
+        if isinstance(node, ast.NameConstant):
+            return KnownValue(node.value)
+        elif isinstance(node, ast.Num):
+            return KnownValue(node.n)
+        elif isinstance(node, (ast.Str, ast.Bytes)):
+            return KnownValue(node.s)
+        else:
+            raise InvalidEvaluation("Only literals supported")
+
+    def get_name(self, node: ast.Name) -> Value:
+        try:
+            return self.ctx.variables[node.id]
+        except KeyError:
+            raise InvalidEvaluation(f"Invalid variable {node.id}") from None
+
+
+@dataclass
+class TypeEvaluator(ast.NodeVisitor):
+    ctx: Context
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> EvalReturn:
+        return self.visit_block(node.body)
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> EvalReturn:
+        return self.visit_block(node.body)
+
+    def visit_block(self, statements: Sequence[ast.stmt]) -> EvalReturn:
+        for stmt in statements:
+            result = self.visit(stmt)
+            if not may_be_none(result):
+                return result
+        return None
+
+    def visit_Return(self, node: ast.Return) -> EvalReturn:
+        if node.value is None:
+            raise InvalidEvaluation("return statement must have a value")
+        return self.ctx.evaluate_type(node.value)
+
+    def visit_Raise(self, node: ast.Raise) -> EvalReturn:
+        if (
+            not isinstance(node.exc, ast.Call)
+            or not isinstance(node.exc.func, ast.Name)
+            or node.exc.func.id != "Exception"
+            or node.exc.keywords
+            or len(node.exc.args) != 1
+            or not isinstance(node.exc.args[0], ast.Str)
+        ):
+            raise InvalidEvaluation(
+                "raise statement must be of the form 'raise Exception(message)'"
+            )
+        return CanAssignError(node.exc.args[0].s)
+
+    def visit_If(self, node: ast.If) -> EvalReturn:
+        visitor = ConditionEvaluator(self.ctx)
+        condition = visitor.visit(node.test)
+        if condition.is_any_match:
+            with self.ctx.narrow_variables(condition.left_varmap):
+                left_result = self.visit_block(node.body)
+            with self.ctx.narrow_variables(condition.right_varmap):
+                right_result = self.visit_block(node.orelse)
+            return AnyCombinedReturn(left_result, right_result)
+        else:
+            if condition.left_varmap is not None:
+                with self.ctx.narrow_variables(condition.left_varmap):
+                    left_result = self.visit_block(node.body)
+            else:
+                left_result = None
+            if condition.right_varmap is not None:
+                with self.ctx.narrow_variables(condition.right_varmap):
+                    right_result = self.visit_block(node.orelse)
+            else:
+                right_result = None
+            if condition.left_varmap is not None:
+                if condition.right_varmap is not None:
+                    return UnionCombinedReturn(left_result, right_result)
+                else:
+                    return left_result
+            else:
+                if condition.right_varmap is not None:
+                    return right_result
+                else:
+                    raise InvalidEvaluation("Condition must either match or not match")
+
+
+def evaluate(evaluator: Evaluator, ctx: Context) -> Union[Value, CanAssignError]:
+    visitor = TypeEvaluator(ctx)
+    try:
+        result = visitor.visit(evaluator.node)
+        return _evaluate_ret(result)
+    except InvalidEvaluation as e:
+        return CanAssignError(
+            "Internal error in type evaluator", [CanAssignError(str(e))]
+        )
+
+
+def _evaluate_ret(ret: EvalReturn) -> Union[Value, CanAssignError]:
+    if ret is None:
+        raise InvalidEvaluation("Evaluator failed to return")
+    elif isinstance(ret, AnyCombinedReturn):
+        left = _evaluate_ret(ret.left)
+        right = _evaluate_ret(ret.right)
+
+        # If one branch is an error, pick the other one
+        if isinstance(left, CanAssignError):
+            if isinstance(right, CanAssignError):
+                return CanAssignError(children=[left, right])
+            return right
+        if isinstance(right, CanAssignError):
+            return left
+
+        if left == right:
+            return left
+        else:
+            return AnyValue(AnySource.multiple_overload_matches)
+    elif isinstance(ret, UnionCombinedReturn):
+        left = _evaluate_ret(ret.left)
+        if isinstance(left, CanAssignError):
+            return left
+        right = _evaluate_ret(ret.right)
+        if isinstance(right, CanAssignError):
+            return right
+        return unite_values(left, right)
+    else:
+        return ret
+
+
+def get_evaluator(func: Callable[..., Any]) -> Optional[Evaluator]:
+    try:
+        key = f"{func.__module__}.{func.__qualname__}"
+    except AttributeError:
+        return None
+    evaluation_func = get_type_evaluation(key)
+    if evaluation_func is None or not hasattr(evaluation_func, "__globals__"):
+        return None
+    lines, _ = inspect.getsourcelines(evaluation_func)
+    code = textwrap.dedent("".join(lines))
+    body = ast.parse(code)
+    if not body.body:
+        return None
+    evaluator = body.body[0]
+    if not isinstance(evaluator, ast.FunctionDef):
+        raise InvalidEvaluation(f"Cannot locate {func}")
+    return Evaluator(evaluation_func, evaluator, evaluation_func.__globals__)
+
+
+def decompose_union(
+    expected_type: Value, parent_value: Value, ctx: CanAssignContext
+) -> Optional[Tuple[TypeVarMap, bool, Value]]:
+    value = unannotate(parent_value)
+    if isinstance(value, MultiValuedValue):
+        tv_maps = []
+        remaining_values = []
+        union_used_any = False
+        for val in value.vals:
+            can_assign, subval_used_any = can_assign_and_used_any(
+                expected_type, val, ctx
+            )
+            if isinstance(can_assign, CanAssignError):
+                remaining_values.append(val)
+            else:
+                if subval_used_any:
+                    union_used_any = True
+                tv_maps.append(can_assign)
+        if tv_maps:
+            tv_map = unify_typevar_maps(tv_maps)
+            assert (
+                remaining_values
+            ), f"all union members matched between {expected_type} and {parent_value}"
+            return tv_map, union_used_any, unite_values(*remaining_values)
+    return None
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index 6604192..174009a 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -1953,6 +1953,12 @@ def unannotate_value(
     return origin, []
 
 
+def unannotate(value: Value) -> Value:
+    if isinstance(value, AnnotatedValue):
+        return value.value
+    return value
+
+
 def unite_and_simplify(*values: Value, limit: int) -> Value:
     united = unite_values(*values)
     if not isinstance(united, MultiValuedValue) or len(united.vals) < limit:
@@ -2291,3 +2297,12 @@ def stringify_object(obj: Any) -> str:
             return f"{obj.__module__}.{obj.__name__}"
     except Exception:
         return repr(obj)
+
+
+def can_assign_and_used_any(
+    param_typ: Value, var_value: Value, ctx: CanAssignContext
+) -> Tuple[CanAssign, bool]:
+    with ctx.reset_any_used():
+        tv_map = param_typ.can_assign(var_value, ctx)
+        used_any = ctx.has_used_any_match()
+    return tv_map, used_any


Test patch:
diff --git a/pyanalyze/test_type_evaluation.py b/pyanalyze/test_type_evaluation.py
new file mode 100644
index 0000000..3346b8f
--- /dev/null
+++ b/pyanalyze/test_type_evaluation.py
@@ -0,0 +1,147 @@
+# static analysis: ignore
+from typing_extensions import Literal
+
+from .value import AnySource, AnyValue, TypedValue, assert_is_value
+from .test_node_visitor import assert_passes
+from .test_name_check_visitor import TestNameCheckVisitorBase
+from .extensions import evaluated, is_set
+
+from typing import Union
+
+
+# These functions must be defined globally because the
+# mechanism for finding the source code doesn't work inside
+# tests.
+@evaluated
+def simple_evaluated(x: int, y: str = ""):
+    if is_set(y):
+        return int
+    else:
+        return str
+
+
+def simple_evaluated(*args: object) -> Union[int, str]:
+    if len(args) >= 2:
+        return 1
+    else:
+        return "x"
+
+
+@evaluated
+def isinstance_evaluated(x: int):
+    if isinstance(x, Literal[1]):
+        return str
+    else:
+        return int
+
+
+def isinstance_evaluated(x: int) -> Union[int, str]:
+    if x == 1:
+        return ""
+    else:
+        return 0
+
+
+@evaluated
+def not_evaluated(x: int):
+    if not isinstance(x, Literal[1]):
+        return str
+    else:
+        return int
+
+
+def not_evaluated(x: int) -> Union[int, str]:
+    if x != 1:
+        return ""
+    else:
+        return 0
+
+
+@evaluated
+def compare_evaluated(x: object):
+    if x is None:
+        return str
+    elif x == 1:
+        return float
+    else:
+        return int
+
+
+def compare_evaluated(x: object) -> Union[int, str, float]:
+    raise NotImplementedError
+
+
+@evaluated
+def nonempty_please(x: str) -> int:
+    if x == "":
+        raise Exception("Non-empty string expected")
+    else:
+        return int
+
+
+def nonempty_please(x: str) -> int:
+    assert x
+    return len(x)
+
+
+class TestTypeEvaluation(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_is_set(self):
+        from pyanalyze.test_type_evaluation import simple_evaluated
+
+        def capybara():
+            assert_is_value(simple_evaluated(1), TypedValue(str))
+            assert_is_value(simple_evaluated(1, "1"), TypedValue(int))
+
+    @assert_passes()
+    def test_isinstance(self):
+        from pyanalyze.test_type_evaluation import isinstance_evaluated
+
+        def capybara(unannotated):
+            assert_is_value(isinstance_evaluated(1), TypedValue(str))
+            assert_is_value(isinstance_evaluated(2), TypedValue(int))
+            assert_is_value(
+                isinstance_evaluated(unannotated),
+                AnyValue(AnySource.multiple_overload_matches),
+            )
+            assert_is_value(
+                isinstance_evaluated(2 if unannotated else 1),
+                TypedValue(int) | TypedValue(str),
+            )
+
+    @assert_passes()
+    def test_not(self):
+        from pyanalyze.test_type_evaluation import not_evaluated
+
+        def capybara(unannotated):
+            assert_is_value(not_evaluated(1), TypedValue(int))
+            assert_is_value(not_evaluated(2), TypedValue(str))
+            assert_is_value(
+                not_evaluated(unannotated),
+                AnyValue(AnySource.multiple_overload_matches),
+            )
+            assert_is_value(
+                not_evaluated(2 if unannotated else 1),
+                TypedValue(int) | TypedValue(str),
+            )
+
+    @assert_passes()
+    def test_compare(self):
+        from pyanalyze.test_type_evaluation import compare_evaluated
+
+        def capybara(unannotated):
+            assert_is_value(compare_evaluated(None), TypedValue(str))
+            assert_is_value(compare_evaluated(1), TypedValue(float))
+            assert_is_value(compare_evaluated("x"), TypedValue(int))
+            assert_is_value(
+                compare_evaluated(None if unannotated else 1),
+                TypedValue(str) | TypedValue(float),
+            )
+
+    @assert_passes()
+    def test_error(self):
+        from pyanalyze.test_type_evaluation import nonempty_please
+
+        def capybara():
+            nonempty_please("")  # E: incompatible_call
+            assert_is_value(nonempty_please("x"), TypedValue(int))