Repository: quora/pyanalyze
Pull request #532 (2022-04-26 23:52:19)

Improve typevar solver for constrained TypeVars


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 63be44f..b182fcb 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Improve `TypeVar` solution heuristic for constrained
+  typevars with multiple solutions (#532)
 - Fix resolution of stringified annotations in `__init__` methods (#530)
 - Type check `yield`, `yield from`, and `return` nodes in generators (#529)
 - Type check calls to comparison operators (#527)
diff --git a/pyanalyze/typevar.py b/pyanalyze/typevar.py
index f3cae50..0ffcb6d 100644
--- a/pyanalyze/typevar.py
+++ b/pyanalyze/typevar.py
@@ -114,8 +114,36 @@ def solve(
             for option, can_assign in zip(options, can_assigns)
             if not isinstance(can_assign, CanAssignError)
         ]
+        # If there's only one solution, pick it.
         if len(available) == 1:
             return available[0]
-        # TODO consider returning unite_values(*available) here instead.
+        # If we inferred Any, keep it; all the solutions will be valid, and
+        # picking one will lead to weird errors down the line.
+        if isinstance(solution, AnyValue):
+            return solution
+        available = remove_redundant_solutions(available, ctx)
+        if len(available) == 1:
+            return available[0]
+        # If there are still multiple options, we fall back to Any.
         return AnyValue(AnySource.inference)
     return solution
+
+
+def remove_redundant_solutions(
+    solutions: Sequence[Value], ctx: CanAssignContext
+) -> Sequence[Value]:
+    # This is going to be quadratic, so don't do it when there's too many
+    # opttions.
+    initial_count = len(solutions)
+    if initial_count > 10:
+        return solutions
+
+    temp_solutions = list(solutions)
+    for i in range(initial_count):
+        sol = temp_solutions[i]
+        for j, other in enumerate(temp_solutions):
+            if i == j or other is None:
+                continue
+            if sol.is_assignable(other, ctx) and not other.is_assignable(sol, ctx):
+                temp_solutions[i] = None
+    return [sol for sol in temp_solutions if sol is not None]


Test patch:
diff --git a/pyanalyze/test_annotations.py b/pyanalyze/test_annotations.py
index febe09e..cb2b50b 100644
--- a/pyanalyze/test_annotations.py
+++ b/pyanalyze/test_annotations.py
@@ -904,9 +904,10 @@ class TestTypeVar(TestNameCheckVisitorBase):
             assert_is_value(f(b"x"), TypedValue(bytes))
             assert_is_value(f(s), TypedValue(str))
             assert_is_value(f(b), TypedValue(bytes))
-            f(sb)  # E: incompatible_argument
+            result = f(sb)  # E: incompatible_argument
+            assert_is_value(result, AnyValue(AnySource.error))
             f(3)  # E: incompatible_argument
-            assert_is_value(f(unannotated), AnyValue(AnySource.inference))
+            assert_is_value(f(unannotated), AnyValue(AnySource.unannotated))
 
     @assert_passes()
     def test_constraint_in_typeshed(self):
diff --git a/pyanalyze/test_typeshed.py b/pyanalyze/test_typeshed.py
index 9df4641..8d01adc 100644
--- a/pyanalyze/test_typeshed.py
+++ b/pyanalyze/test_typeshed.py
@@ -667,9 +667,9 @@ class TestParamSpec(TestNameCheckVisitorBase):
             wrapped("x")  # E: incompatible_argument
 
 
-class TestOpen(TestNameCheckVisitorBase):
+class TestIntegration(TestNameCheckVisitorBase):
     @assert_passes()
-    def test_basic(self):
+    def test_open(self):
         import io
         from typing import Any, BinaryIO, IO
 
@@ -682,3 +682,12 @@ class TestOpen(TestNameCheckVisitorBase):
             assert_type(open("x", "rb", buffering=0), io.FileIO)
             assert_type(open("x", "rb", buffering=buffering), BinaryIO)
             assert_type(open("x", mode, buffering=buffering), IO[Any])
+
+    @assert_passes()
+    def test_itertools_count(self):
+        import itertools
+
+        def capybara():
+            assert_is_value(
+                itertools.count(1), GenericValue(itertools.count, [TypedValue(int)])
+            )
diff --git a/pyanalyze/test_typevar.py b/pyanalyze/test_typevar.py
index b559e2b..86e0d66 100644
--- a/pyanalyze/test_typevar.py
+++ b/pyanalyze/test_typevar.py
@@ -6,6 +6,7 @@ from .value import (
     AnnotatedValue,
     AnySource,
     AnyValue,
+    GenericValue,
     KnownValue,
     MultiValuedValue,
     TypedValue,
@@ -260,6 +261,34 @@ class TestSolve(TestNameCheckVisitorBase):
             m = min(E)
             assert_is_value(m, TypedValue(E))
 
+    @assert_passes()
+    def test_constraints(self):
+        from typing import List, TypeVar
+
+        LT = TypeVar("LT", List[int], List[str])
+
+        def g(x: LT) -> LT:
+            return x
+
+        def pacarana() -> None:
+            assert_is_value(g([]), AnyValue(AnySource.inference))
+            assert_is_value(g([1]), GenericValue(list, [TypedValue(int)]))
+
+    @assert_passes()
+    def test_redundant_constraints(self):
+        from typing import TypeVar
+        from typing_extensions import SupportsIndex
+
+        T = TypeVar("T", int, float, SupportsIndex)
+
+        def f(x: T) -> T:
+            return x
+
+        def capybara(si: SupportsIndex):
+            assert_is_value(f(1), TypedValue(int))
+            assert_is_value(f(si), TypedValue(SupportsIndex))
+            assert_is_value(f(1.0), TypedValue(float))
+
 
 class TestAnnotated(TestNameCheckVisitorBase):
     @assert_passes()