Repository: quora/pyanalyze
Pull request #409 (2022-01-12 13:46:01)

support generic type eval


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 5a0e681..1156a92 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Support generic type evaluators (#409)
 - Implement return annotation behavior for type evaluation
   functions (#408)
 - Support `extend_config` option in `pyproject.toml` (#407)
diff --git a/docs/type_evaluation.md b/docs/type_evaluation.md
index eb2be40..d33f4b4 100644
--- a/docs/type_evaluation.md
+++ b/docs/type_evaluation.md
@@ -37,6 +37,13 @@ Type evaluation functions can replace most complex overloads
 with simpler, more readable code. They solve a number of
 problems:
 
+- 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`.
+  - Checking whether two generic arguments are overlapping
 - 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
@@ -48,12 +55,6 @@ problems:
 - 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
 
@@ -403,6 +404,31 @@ For example:
     union: int | str
     reveal_type(switch_types(union))  # int | str
 
+### Generic evaluators
+
+If any type variables appear in the parameters of the type evaluation
+function, the type checker should first solve those and use the solution
+in the body of the function:
+
+    @evaluated
+    def identity(x: T):
+        return T
+
+    reveal_type(evaluated(int()))  # int
+
+As a result, `is_of_type()` checks that use a type variable work:
+
+    @evaluated
+    def safe_upcast(typ: Type[T1], value: object):
+        if is_of_type(value, T1):
+            return T1
+        show_error("unsafe cast")
+        return Any
+
+    reveal_type(safe_upcast(object, 1))  # object
+    reveal_type(safe_upcast(int, 1))  # int
+    safe_upcast(str, 1)  # error
+
 ### Type compatibility
 
 The type of an evaluated function is compatible with a
@@ -587,6 +613,28 @@ 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
 
+### Generic evaluators
+
+The specification for generic evaluators allows creating an evaluator
+that checks whether two types have any overlap:
+
+    T1 = TypeVar("T1")
+    T2 = TypeVar("T2")
+
+    @evaluated
+    def safe_contains(elt: T1, container: Container[T2]) -> bool:
+        if not is_of_type(elt, T2) and not is_of_type(container, Container[T1]):
+            show_error("Element cannot be a member of container")
+
+    lst: List[int]
+    safe_contains("x", lst)  # error
+    safe_contains(True, lst)  # ok (bool is a subclass of int)
+    safe_contains(object(), lst)  # ok (List[int] is a subclass of Container[object])
+
+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.
+
 ## Possible extensions
 
 The following features may be useful, but are deferred
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 1f00e18..d128eef 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -1813,6 +1813,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                     variables={param.name: param.annotation for param in params},
                     positions={param.name: type_evaluation.DEFAULT for param in params},
                     can_assign_context=self,
+                    tv_map={},
                 )
                 for error in evaluator.validate(ctx):
                     self.show_error(
diff --git a/pyanalyze/signature.py b/pyanalyze/signature.py
index c054508..d445af8 100644
--- a/pyanalyze/signature.py
+++ b/pyanalyze/signature.py
@@ -1052,7 +1052,7 @@ class Signature:
             ctx.visitor.record_call(self.callable, variables)
 
         return_value = self.return_value
-        typevar_values: Dict[TypeVar, Value] = {}
+        typevar_values: TypeVarMap = {}
         if self.all_typevars:
             tv_possible_values: Dict[TypeVar, List[Value]] = defaultdict(list)
             for param_name in self.typevars_of_params:
@@ -1140,7 +1140,9 @@ class Signature:
                 positions = {
                     param: position for param, (position, _) in bound_args.items()
                 }
-                eval_ctx = EvalContext(varmap, positions, ctx.can_assign_ctx)
+                eval_ctx = EvalContext(
+                    varmap, positions, ctx.can_assign_ctx, typevar_values
+                )
                 return_value, errors = self.evaluator.evaluate(eval_ctx)
                 for error in errors:
                     error_node = None
diff --git a/pyanalyze/type_evaluation.py b/pyanalyze/type_evaluation.py
index f91ee84..ff750db 100644
--- a/pyanalyze/type_evaluation.py
+++ b/pyanalyze/type_evaluation.py
@@ -261,6 +261,7 @@ class EvalContext:
     variables: VarMap
     positions: Mapping[str, Position]
     can_assign_context: CanAssignContext
+    tv_map: TypeVarMap
 
     @contextmanager
     def narrow_variables(self, varmap: Optional[VarMap]) -> Iterator[None]:
@@ -300,6 +301,10 @@ class Evaluator:
     def evaluate_value(self, __node: ast.AST) -> Value:
         raise NotImplementedError
 
+    def evaluate_generic_type(self, __node: ast.AST, __ctx: EvalContext) -> Value:
+        typ = self.evaluate_type(__node)
+        return typ.substitute_typevars(__ctx.tv_map)
+
 
 @dataclass
 class CombinedReturn:
@@ -387,7 +392,7 @@ class ConditionEvaluator(ast.NodeVisitor):
                     "is_of_type() takes two positional arguments", node
                 )
             varname_node = node.args[0]
-            typ = self.evaluator.evaluate_type(node.args[1])
+            typ = self.evaluator.evaluate_generic_type(node.args[1], self.ctx)
             exclude_any = True
             for keyword in node.keywords:
                 if keyword.arg == "exclude_any":
@@ -628,7 +633,7 @@ class EvaluateVisitor(ast.NodeVisitor):
 
     def _evaluate_ret(self, ret: EvalReturn, node: ast.AST) -> Value:
         if ret is None:
-            return self.evaluator.return_annotation
+            return self.evaluator.return_annotation.substitute_typevars(self.ctx.tv_map)
         elif isinstance(ret, CombinedReturn):
             children = [self._evaluate_ret(child, node) for child in ret.children]
             return unite_values(*children)
@@ -678,7 +683,7 @@ class EvaluateVisitor(ast.NodeVisitor):
         if node.value is None:
             self.add_invalid("return statement must have a value", node)
             return KnownValue(None)
-        return self.evaluator.evaluate_type(node.value)
+        return self.evaluator.evaluate_generic_type(node.value, self.ctx)
 
     def visit_Expr(self, node: ast.Expr) -> EvalReturn:
         if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name):


Test patch:
diff --git a/pyanalyze/test_type_evaluation.py b/pyanalyze/test_type_evaluation.py
index b6ba9ea..34d6634 100644
--- a/pyanalyze/test_type_evaluation.py
+++ b/pyanalyze/test_type_evaluation.py
@@ -1,5 +1,5 @@
 # static analysis: ignore
-from .value import KnownValue, TypedValue, assert_is_value
+from .value import AnySource, AnyValue, KnownValue, TypedValue, assert_is_value
 from .test_node_visitor import assert_passes
 from .test_name_check_visitor import TestNameCheckVisitorBase
 from .extensions import is_keyword, is_positional, is_provided, is_of_type, show_error
@@ -309,6 +309,27 @@ class TestTypeEvaluation(TestNameCheckVisitorBase):
             assert_is_value(maybe_use_header(True), TypedValue(str))
             assert_is_value(maybe_use_header(x), TypedValue(int))
 
+    @assert_passes()
+    def test_generic(self):
+        from pyanalyze.extensions import evaluated
+        from typing import TypeVar
+
+        T1 = TypeVar("T1")
+
+        @evaluated
+        def identity(x: T1):
+            return T1
+
+        @evaluated
+        def identity2(x: T1) -> T1:
+            pass
+
+        def capybara(unannotated):
+            assert_is_value(identity(1), KnownValue(1))
+            assert_is_value(identity(unannotated), AnyValue(AnySource.unannotated))
+            assert_is_value(identity2(1), KnownValue(1))
+            assert_is_value(identity2(unannotated), AnyValue(AnySource.unannotated))
+
 
 class TestBoolOp(TestNameCheckVisitorBase):
     @assert_passes()
@@ -508,3 +529,41 @@ class TestExamples(TestNameCheckVisitorBase):
             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))
+
+    @assert_passes()
+    def test_safe_upcast(self):
+        from typing import Type, Any, TypeVar
+        from pyanalyze.extensions import evaluated, show_error, is_of_type
+
+        T1 = TypeVar("T1")
+
+        @evaluated
+        def safe_upcast(typ: Type[T1], value: object):
+            if is_of_type(value, T1):
+                return T1
+            show_error("unsafe cast")
+            return Any
+
+        def capybara():
+            assert_is_value(safe_upcast(object, 1), TypedValue(object))
+            assert_is_value(safe_upcast(int, 1), TypedValue(int))
+            safe_upcast(str, 1)  # E: incompatible_call
+
+    @assert_passes()
+    def test_safe_contains(self):
+        from typing import List, TypeVar, Container
+        from pyanalyze.extensions import evaluated, show_error, is_of_type
+
+        T1 = TypeVar("T1")
+        T2 = TypeVar("T2")
+
+        @evaluated
+        def safe_contains(elt: T1, container: Container[T2]) -> bool:
+            if not is_of_type(elt, T2) and not is_of_type(container, Container[T1]):
+                show_error("Element cannot be a member of container")
+
+        def capybara(lst: List[int], o: object):
+            safe_contains(True, ["x"])  # E: incompatible_call
+            safe_contains("x", lst)  # E: incompatible_call
+            safe_contains(True, lst)
+            safe_contains(o, lst)