Repository: quora/pyanalyze
Pull request #408 (2022-01-12 13:11:32)

implement return annotation behavior for type eval


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 22a7ff3..5a0e681 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Implement return annotation behavior for type evaluation
+  functions (#408)
 - Support `extend_config` option in `pyproject.toml` (#407)
 - Remove the old method return type check. Use the new
   `incompatible_override` check instead (#404)
diff --git a/docs/type_evaluation.md b/docs/type_evaluation.md
index 2706805..eb2be40 100644
--- a/docs/type_evaluation.md
+++ b/docs/type_evaluation.md
@@ -690,13 +690,10 @@ in pyanalyze:
 
 Currently unsupported features include:
 
-- Usage in stubs
 - pyanalyze should provide a way to register
   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:
 
diff --git a/pyanalyze/annotations.py b/pyanalyze/annotations.py
index f2d15fd..f189e61 100644
--- a/pyanalyze/annotations.py
+++ b/pyanalyze/annotations.py
@@ -30,9 +30,7 @@ import typing_inspect
 import qcore
 import ast
 import builtins
-import inspect
 from collections.abc import Callable, Iterable
-import textwrap
 from typing import (
     Any,
     Container,
@@ -59,7 +57,6 @@ from .extensions import (
     NoReturnGuard,
     ParameterTypeGuard,
     TypeGuard,
-    get_type_evaluation,
 )
 from .find_unused import used
 from .functions import FunctionDefNode
@@ -170,25 +167,6 @@ class RuntimeEvaluator(type_evaluation.Evaluator, Context):
         """Return the :class:`Value <pyanalyze.value.Value>` corresponding to a name."""
         return self.get_name_from_globals(node.id, self.globals)
 
-    @classmethod
-    def get_for(cls, func: typing.Callable[..., Any]) -> Optional["RuntimeEvaluator"]:
-        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):
-            return None
-        return RuntimeEvaluator(evaluator, evaluation_func.__globals__, evaluation_func)
-
 
 @dataclass
 class SyntheticEvaluator(type_evaluation.Evaluator):
@@ -217,10 +195,16 @@ class SyntheticEvaluator(type_evaluation.Evaluator):
 
     @classmethod
     def from_visitor(
-        cls, node: FunctionDefNode, visitor: "NameCheckVisitor"
+        cls,
+        node: FunctionDefNode,
+        visitor: "NameCheckVisitor",
+        return_annotation: Value,
     ) -> "SyntheticEvaluator":
         return cls(
-            node, visitor, _DefaultContext(visitor, node, use_name_node_for_error=True)
+            node,
+            return_annotation,
+            visitor,
+            _DefaultContext(visitor, node, use_name_node_for_error=True),
         )
 
 
diff --git a/pyanalyze/arg_spec.py b/pyanalyze/arg_spec.py
index 3d5a0c6..0e2e8cd 100644
--- a/pyanalyze/arg_spec.py
+++ b/pyanalyze/arg_spec.py
@@ -6,7 +6,7 @@ Implementation of extended argument specifications used by test_scope.
 
 from .options import Options, PyObjectSequenceOption
 from .analysis_lib import is_positional_only_arg_name
-from .extensions import CustomCheck, get_overloads
+from .extensions import CustomCheck, get_overloads, get_type_evaluation
 from .annotations import Context, RuntimeEvaluator, type_from_runtime
 from .config import Config
 from .find_unused import used
@@ -60,6 +60,7 @@ from dataclasses import dataclass, replace
 import qcore
 import inspect
 import sys
+import textwrap
 from types import FunctionType, ModuleType
 from typing import (
     Any,
@@ -509,6 +510,37 @@ class ArgSpecCache:
             self.known_argspecs[obj] = extended
         return extended
 
+    def _maybe_make_evaluator_sig(
+        self, func: Callable[..., Any], impl: Optional[Impl], is_asynq: bool
+    ) -> MaybeSignature:
+        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
+        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)
+
     def _uncached_get_argspec(
         self,
         obj: Any,
@@ -529,13 +561,9 @@ class ArgSpecCache:
                     ]
                     if all_of_type(sigs, Signature):
                         return OverloadedSignature(sigs)
-                evaluator = RuntimeEvaluator.get_for(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)
+                evaluator_sig = self._maybe_make_evaluator_sig(obj, impl, is_asynq)
+                if evaluator_sig is not None:
+                    return evaluator_sig
 
         if isinstance(obj, tuple) or hasattr(obj, "__getattr__"):
             return None  # lost cause
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index f911ed0..1f00e18 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -1805,7 +1805,10 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             if self._is_collecting() or isinstance(node, ast.Lambda):
                 return FunctionResult(parameters=params)
             with self.scopes.allow_only_module_scope():
-                evaluator = SyntheticEvaluator.from_visitor(node, self)
+                # The return annotation doesn't actually matter for validation.
+                evaluator = SyntheticEvaluator.from_visitor(
+                    node, self, AnyValue(AnySource.marker)
+                )
                 ctx = type_evaluation.EvalContext(
                     variables={param.name: param.annotation for param in params},
                     positions={param.name: type_evaluation.DEFAULT for param in params},
diff --git a/pyanalyze/type_evaluation.py b/pyanalyze/type_evaluation.py
index 4d5348d..f91ee84 100644
--- a/pyanalyze/type_evaluation.py
+++ b/pyanalyze/type_evaluation.py
@@ -36,8 +36,6 @@ from .stacked_scopes import (
 from .safe import all_of_type
 from .value import (
     NO_RETURN_VALUE,
-    AnySource,
-    AnyValue,
     CanAssign,
     CanAssignContext,
     CanAssignError,
@@ -281,6 +279,7 @@ class EvalContext:
 @dataclass
 class Evaluator:
     node: Union[ast.FunctionDef, ast.AsyncFunctionDef]
+    return_annotation: Value
 
     def evaluate(self, ctx: EvalContext) -> Tuple[Value, Sequence[UserRaisedError]]:
         visitor = EvaluateVisitor(self, ctx)
@@ -629,10 +628,7 @@ class EvaluateVisitor(ast.NodeVisitor):
 
     def _evaluate_ret(self, ret: EvalReturn, node: ast.AST) -> Value:
         if ret is None:
-            # TODO return the func's return annotation instead
-            if not self.validation_mode:
-                self.add_invalid("Evaluator failed to return", node)
-            return AnyValue(AnySource.error)
+            return self.evaluator.return_annotation
         elif isinstance(ret, CombinedReturn):
             children = [self._evaluate_ret(child, node) for child in ret.children]
             return unite_values(*children)
diff --git a/pyanalyze/typeshed.py b/pyanalyze/typeshed.py
index d95f3b9..2856b75 100644
--- a/pyanalyze/typeshed.py
+++ b/pyanalyze/typeshed.py
@@ -764,7 +764,9 @@ class TypeshedFinder:
             cleaned_arguments.append(arg)
         if is_evaluated:
             ctx = _AnnotationContext(self, mod)
-            evaluator = SyntheticEvaluator(node, _DummyErrorContext(), ctx)
+            evaluator = SyntheticEvaluator(
+                node, return_value, _DummyErrorContext(), ctx
+            )
         else:
             evaluator = None
         return Signature.make(


Test patch:
diff --git a/pyanalyze/stubs/_pyanalyze_tests-stubs/evaluated.pyi b/pyanalyze/stubs/_pyanalyze_tests-stubs/evaluated.pyi
index 1ea1d34..def401f 100644
--- a/pyanalyze/stubs/_pyanalyze_tests-stubs/evaluated.pyi
+++ b/pyanalyze/stubs/_pyanalyze_tests-stubs/evaluated.pyi
@@ -9,3 +9,10 @@ def open(mode: str):
         return BinaryIO
     else:
         return IO[Any]
+
+@evaluated
+def open2(mode: str) -> IO[Any]:
+    if mode == "r":
+        return TextIO
+    elif mode == "rb":
+        return BinaryIO
diff --git a/pyanalyze/test_type_evaluation.py b/pyanalyze/test_type_evaluation.py
index 0763bb9..b6ba9ea 100644
--- a/pyanalyze/test_type_evaluation.py
+++ b/pyanalyze/test_type_evaluation.py
@@ -296,6 +296,19 @@ class TestTypeEvaluation(TestNameCheckVisitorBase):
             has_default()  # E: incompatible_call
             has_default(i)  # E: incompatible_call
 
+    @assert_passes()
+    def test_return(self):
+        from pyanalyze.extensions import evaluated
+
+        @evaluated
+        def maybe_use_header(x: bool) -> int:
+            if x is True:
+                return str
+
+        def capybara(x: bool):
+            assert_is_value(maybe_use_header(True), TypedValue(str))
+            assert_is_value(maybe_use_header(x), TypedValue(int))
+
 
 class TestBoolOp(TestNameCheckVisitorBase):
     @assert_passes()
diff --git a/pyanalyze/test_typeshed.py b/pyanalyze/test_typeshed.py
index 1774a15..809607c 100644
--- a/pyanalyze/test_typeshed.py
+++ b/pyanalyze/test_typeshed.py
@@ -223,7 +223,7 @@ class TestBundledStubs(TestNameCheckVisitorBase):
     @assert_passes()
     def test_evaluated_import(self):
         def capybara(unannotated):
-            from _pyanalyze_tests.evaluated import open
+            from _pyanalyze_tests.evaluated import open, open2
             from typing import TextIO, BinaryIO, IO
 
             assert_is_value(open("r"), TypedValue(TextIO))
@@ -235,6 +235,15 @@ class TestBundledStubs(TestNameCheckVisitorBase):
                 open("r" if unannotated else "rb"),
                 TypedValue(TextIO) | TypedValue(BinaryIO),
             )
+            assert_is_value(open2("r"), TypedValue(TextIO))
+            assert_is_value(open2("rb"), TypedValue(BinaryIO))
+            assert_is_value(
+                open2(unannotated), GenericValue(IO, [AnyValue(AnySource.explicit)])
+            )
+            assert_is_value(
+                open2("r" if unannotated else "rb"),
+                TypedValue(TextIO) | TypedValue(BinaryIO),
+            )
 
     @assert_passes()
     def test_recursive_base(self):