type evaluation docs (#412)

Author: JelleZijlstra

docs/changelog.md

@@ -50,7 +50,7 @@
 - Support the Python 3.10 `match` statement (#376)
 - Support the walrus (`:=`) operator (#375)
 - Initial support for proposed new "type evaluation"
-  mechanism (#374, #379, #384)
+  mechanism (#374, #379, #384, #410)
 - Create command-line options for each config option (#373)
 - Overhaul treatment of function definitions (#372)
   - Support positional-only arguments

docs/type_evaluation.md

@@ -339,7 +339,7 @@ Examples:
 
     @evaluated
     def length_or_none(s: str | None = None):
-        if is_of_type(s, str):
+        if is_of_type(s, str, exclude_any=False):
             return int
         else:
             return None
@@ -353,9 +353,9 @@ Examples:
 
     @evaluated
     def length_or_none2(s: str | None):
-        if is_of_type(s, str, exclude_any=True):
+        if is_of_type(s, str):
             return int
-        elif is_of_type(s, None, exclude_any=True):
+        elif is_of_type(s, None):
             return None
         else:
             return Any
@@ -367,9 +367,9 @@ Examples:
 
     @evaluated
     def nested_any(s: Sequence[Any]):
-        if is_of_type(s, str, exclude_any=True):
+        if is_of_type(s, str):
             show_error("error")
-        elif is_of_type(s, Sequence[str], exclude_any=True):
+        elif is_of_type(s, Sequence[str]):
             return str
         else:
             return int
@@ -453,6 +453,31 @@ Examples:
     _: Callable[[str], Path | None] = maybe_path  # ok
     _: Callable[[Literal["x"]], Path] = maybe_path  # ok
 
+### Runtime behavior
+
+At runtime, the `@evaluated` decorator returns a dummy function
+that throws an error when called, similar to `@overload`. In
+order to support dynamic type checkers, it also stores the
+original function, keyed by its fully qualified name.
+
+A helper function is provided to retrieve all registered
+evaluation functions for a given fully qualified name:
+
+    def get_type_evaluations(
+        fully_qualified_name: str
+    ) -> Sequence[Callable[..., Any]]: ...
+
+For example, if method `B.c` in module `a` has an evaluation function,
+`get_type_evaluations("a.B.c")` will retrieve it.
+
+Dummy implementations are provided for the various helper
+functions (`is_provided()`, `is_positional()`, `is_keyword()`,
+`is_of_type()`, and `show_error()`). These throw an error
+if called at runtime.
+
+The `reveal_type()` function has a runtime implementation
+that simply returns its argument.
+
 ## Discussion
 
 ### Interaction with Any
@@ -635,6 +660,17 @@ Thus, type evaluation provides a way to implement checks similar to mypy's
 [strict equality](https://mypy.readthedocs.io/en/stable/command_line.html#cmdoption-mypy-strict-equality)
 flag directly in stubs.
 
+## Compatibility
+
+The proposal is fully backward compatible.
+
+Type evaluation functions are going to be most frequently useful
+in library stubs, where it is often important that multiple type
+checkers can parse the stub. In order to unblock usage of the new
+feature in stubs, type checker authors could simply ignore the
+body of evaluation functions and rely on the signature. This would
+still allow other type checkers to fully use the evaluation function.
+
 ## Possible extensions
 
 The following features may be useful, but are deferred

pyanalyze/arg_spec.py

@@ -6,7 +6,7 @@
 
 from .options import Options, PyObjectSequenceOption
 from .analysis_lib import is_positional_only_arg_name
-from .extensions import CustomCheck, get_overloads, get_type_evaluation
+from .extensions import CustomCheck, get_overloads, get_type_evaluations
 from .annotations import Context, RuntimeEvaluator, type_from_runtime
 from .config import Config
 from .find_unused import used
@@ -517,29 +517,36 @@ def _maybe_make_evaluator_sig(
             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__"):
+        evaluation_funcs = get_type_evaluations(key)
+        if not evaluation_funcs:
             return None
-        sig = self._cached_get_argspec(
-            evaluation_func, impl, is_asynq, in_overload_resolution=True
-        )
-        if sig is None:
-            return None
-        lines, _ = inspect.getsourcelines(evaluation_func)
-        code = textwrap.dedent("".join(lines))
-        body = ast.parse(code)
-        if not body.body:
-            return None
-        evaluator_node = body.body[0]
-        if not isinstance(evaluator_node, ast.FunctionDef):
-            return None
-        evaluator = RuntimeEvaluator(
-            evaluator_node,
-            sig.return_value,
-            evaluation_func.__globals__,
-            evaluation_func,
-        )
-        return replace(sig, evaluator=evaluator)
+        sigs = []
+        for evaluation_func in evaluation_funcs:
+            if evaluation_func is None or not hasattr(evaluation_func, "__globals__"):
+                return None
+            sig = self._cached_get_argspec(
+                evaluation_func, impl, is_asynq, in_overload_resolution=True
+            )
+            if not isinstance(sig, Signature):
+                return None
+            lines, _ = inspect.getsourcelines(evaluation_func)
+            code = textwrap.dedent("".join(lines))
+            body = ast.parse(code)
+            if not body.body:
+                return None
+            evaluator_node = body.body[0]
+            if not isinstance(evaluator_node, ast.FunctionDef):
+                return None
+            evaluator = RuntimeEvaluator(
+                evaluator_node,
+                sig.return_value,
+                evaluation_func.__globals__,
+                evaluation_func,
+            )
+            sigs.append(replace(sig, evaluator=evaluator))
+        if len(sigs) == 1:
+            return sigs[0]
+        return OverloadedSignature(sigs)
 
     def _uncached_get_argspec(
         self,

pyanalyze/extensions.py

@@ -19,6 +19,7 @@
     Dict,
     Iterable,
     Optional,
+    Sequence,
     Tuple,
     List,
     Union,
@@ -387,17 +388,17 @@ def f(x: int) -> None:
 
 
 _overloads: Dict[str, List[Callable[..., Any]]] = defaultdict(list)
-_type_evaluations: Dict[str, Optional[Callable[..., Any]]] = {}
+_type_evaluations: Dict[str, List[Callable[..., Any]]] = defaultdict(list)
 
 
 def get_overloads(fully_qualified_name: str) -> List[Callable[..., Any]]:
     """Return all defined runtime overloads for this fully qualified name."""
     return _overloads[fully_qualified_name]
 
 
-def get_type_evaluation(fully_qualified_name: str) -> Optional[Callable[..., Any]]:
+def get_type_evaluations(fully_qualified_name: str) -> Sequence[Callable[..., Any]]:
     """Return the type evaluation function for this fully qualified name, or None."""
-    return _type_evaluations.get(fully_qualified_name)
+    return _type_evaluations[fully_qualified_name]
 
 
 if TYPE_CHECKING:
@@ -433,8 +434,7 @@ def patch_typing_overload() -> None:
 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
+    _type_evaluations[key].append(func)
     func.__is_type_evaluation__ = True
     return func
 

pyanalyze/stubs/pyanalyze-stubs/extensions.pyi

@@ -2,11 +2,11 @@
 # from it, because typeshed_client doesn't let
 # stubs import from non-stub files.
 
-from typing import Any, Callable, Optional, List
+from typing import Any, Callable, Optional, List, Sequence
 
 def reveal_type(value: object) -> None: ...
 def get_overloads(fully_qualified_name: str) -> List[Callable[..., Any]]: ...
-def get_type_evaluation(fully_qualified_name: str) -> Optional[Callable[..., Any]]: ...
+def get_type_evaluation(fully_qualified_name: str) -> Sequence[Callable[..., Any]]: ...
 def overload(func: Callable[..., Any]) -> Callable[..., Any]: ...
 def evaluated(func: Callable[..., Any]) -> Callable[..., Any]: ...
 def is_provided(arg: Any) -> bool: ...