Repository: quora/pyanalyze (python)

Issue:
Support more TypeVar features
We have some support for typevars now, but currently pyanalyzes ignores most non-standard options:
- [ ] Variance is ignored. I'm inclined to keep treating invariant typevars as covariant (the current behavior), because in my experience, the way the Python type system defaults to invariance is more likely to lead to annoying false positives than to find genuine issues. And this is likely to be even worse with pyanalyze than with other type checkers, because pyanalyze is more likely to infer precise Literal types.
- [ ] Contravariant typevars should be supported though. We'll need some concept of contravariance for Callable types (#52) anyway.
- [ ] Typevars with value restriction
- [ ] Typevars with bounds

Pull request:
Support TypeVar bounds and constraints
Fixes #54

Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index ca12c7d..e966be7 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Support `TypeVar` bounds and constraints (#315)
 - Improve error messages involving concrete dictionary and sequence values (#312)
 - More precise type inference for dict literals (#312)
 - Support `AsynqCallable` with no arguments as an annotation (#314)
diff --git a/pyanalyze/annotations.py b/pyanalyze/annotations.py
index 6555b22..bbc9316 100644
--- a/pyanalyze/annotations.py
+++ b/pyanalyze/annotations.py
@@ -350,7 +350,18 @@ def _type_from_runtime(val: Any, ctx: Context, is_typeddict: bool = False) -> Va
             ctx.show_error(f"Invalid NewType {val}")
             return AnyValue(AnySource.error)
     elif typing_inspect.is_typevar(val):
-        return TypeVarValue(cast(TypeVar, val))
+        tv = cast(TypeVar, val)
+        if tv.__bound__ is not None:
+            bound = _type_from_runtime(tv.__bound__, ctx)
+        else:
+            bound = None
+        if tv.__constraints__:
+            constraints = tuple(
+                _type_from_runtime(constraint, ctx) for constraint in tv.__constraints__
+            )
+        else:
+            constraints = ()
+        return TypeVarValue(tv, bound=bound, constraints=constraints)
     elif is_typing_name(val, "Final") or is_typing_name(val, "ClassVar"):
         return AnyValue(AnySource.incomplete_annotation)
     elif typing_inspect.is_classvar(val) or typing_inspect.is_final_type(val):
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 376ade5..7b101f2 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -3769,9 +3769,15 @@ class NameCheckVisitor(
         Returns :data:`pyanalyze.value.UNINITIALIZED_VALUE` if the attribute cannot be found.
 
         """
+        if isinstance(root_composite.value, TypeVarValue):
+            root_composite = Composite(
+                value=root_composite.value.get_fallback_value(),
+                varname=root_composite.varname,
+                node=root_composite.node,
+            )
         if isinstance(root_composite.value, MultiValuedValue):
             values = [
-                self._get_attribute_no_mvv(
+                self.get_attribute(
                     Composite(subval, root_composite.varname, root_composite.node),
                     attr,
                     node,
@@ -3804,13 +3810,19 @@ class NameCheckVisitor(
         self, root_composite: Composite, attr: str, node: ast.AST
     ) -> Value:
         ignore_none = self.config.IGNORE_NONE_ATTRIBUTES
+        if isinstance(root_composite.value, TypeVarValue):
+            root_composite = Composite(
+                value=root_composite.value.get_fallback_value(),
+                varname=root_composite.varname,
+                node=root_composite.node,
+            )
         if isinstance(root_composite.value, MultiValuedValue):
             results = []
             for subval in root_composite.value.vals:
                 composite = Composite(
                     subval, root_composite.varname, root_composite.node
                 )
-                subresult = self._get_attribute_no_mvv(
+                subresult = self.get_attribute(
                     composite, attr, node, ignore_none=ignore_none
                 )
                 if subresult is UNINITIALIZED_VALUE:
diff --git a/pyanalyze/safe.py b/pyanalyze/safe.py
index 04962bf..1f0929f 100644
--- a/pyanalyze/safe.py
+++ b/pyanalyze/safe.py
@@ -108,7 +108,6 @@ if sys.version_info >= (3, 10):
     def is_newtype(obj: object) -> bool:
         return isinstance(obj, NewType)
 
-
 else:
 
     def is_newtype(obj: object) -> bool:
diff --git a/pyanalyze/signature.py b/pyanalyze/signature.py
index afbce6b..0792368 100644
--- a/pyanalyze/signature.py
+++ b/pyanalyze/signature.py
@@ -876,7 +876,16 @@ class Signature:
                 if param.annotation is EMPTY:
                     continue
                 tv_map = param.annotation.can_assign(var_value, visitor)
-                if not isinstance(tv_map, CanAssignError):
+                if isinstance(tv_map, CanAssignError):
+                    visitor.show_error(
+                        node,
+                        f"Incompatible argument type for {param.name}: expected"
+                        f" {param.annotation} but got {var_value}",
+                        ErrorCode.incompatible_argument,
+                        detail=str(tv_map),
+                    )
+                    return self.get_default_return()
+                else:
                     # For now, the first assignment wins.
                     for typevar, value in tv_map.items():
                         tv_possible_values[typevar].append(value)
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index f54f41c..e8f81ba 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -93,7 +93,7 @@ class Value:
         elif isinstance(other, AnnotatedValue):
             return self.can_assign(other.value, ctx)
         elif isinstance(other, TypeVarValue):
-            return self.can_assign(other.get_fallback_value(), ctx)
+            return other.can_be_assigned(self, ctx)
         elif (
             isinstance(other, UnboundMethodValue)
             and other.secondary_attr_name is not None
@@ -1283,6 +1283,8 @@ class MultiValuedValue(Value):
         )
 
     def can_assign(self, other: Value, ctx: CanAssignContext) -> CanAssign:
+        if isinstance(other, TypeVarValue):
+            other = other.get_fallback_value()
         if isinstance(other, MultiValuedValue):
             tv_maps = []
             for val in other.vals:
@@ -1431,17 +1433,88 @@ class TypeVarValue(Value):
     """
 
     typevar: TypeVar
+    bound: Optional[Value] = None
+    constraints: Sequence[Value] = ()
 
     def substitute_typevars(self, typevars: TypeVarMap) -> Value:
         return typevars.get(self.typevar, self)
 
     def can_assign(self, other: Value, ctx: CanAssignContext) -> CanAssign:
+        if self == other:
+            return {}
+        if self.bound is not None:
+            can_assign = self.bound.can_assign(other, ctx)
+            if isinstance(can_assign, CanAssignError):
+                return CanAssignError(
+                    f"Value of TypeVar {self} cannot be {other}", [can_assign]
+                )
+            return {**can_assign, self.typevar: other}
+        elif self.constraints:
+            can_assigns = [
+                constraint.can_assign(other, ctx) for constraint in self.constraints
+            ]
+            if all_of_type(can_assigns, CanAssignError):
+                return CanAssignError(f"Cannot assign to {self}", list(can_assigns))
+            possibilities = [
+                constraint
+                for constraint, can_assign in zip(self.constraints, can_assigns)
+                if not isinstance(can_assign, CanAssignError)
+            ]
+            if len(possibilities) == 1:
+                (solution,) = possibilities
+            else:
+                # Inferring something else produces too many issues for now.
+                solution = AnyValue(AnySource.inference)
+            return {self.typevar: solution}
         return {self.typevar: other}
 
+    def can_be_assigned(self, left: Value, ctx: CanAssignContext) -> CanAssign:
+        if left == self:
+            return {}
+        if self.bound is not None:
+            can_assign = left.can_assign(self.bound, ctx)
+            if isinstance(can_assign, CanAssignError):
+                return CanAssignError(
+                    f"Value of TypeVar {self} cannot be {left}", [can_assign]
+                )
+            return {**can_assign, self.typevar: left}
+        elif self.constraints:
+            can_assigns = [
+                left.can_assign(constraint, ctx) for constraint in self.constraints
+            ]
+            if all_of_type(can_assigns, CanAssignError):
+                return CanAssignError(f"Cannot assign to {self}", list(can_assigns))
+            possibilities = [
+                constraint
+                for constraint, can_assign in zip(self.constraints, can_assigns)
+                if not isinstance(can_assign, CanAssignError)
+            ]
+            if len(possibilities) == 1:
+                (solution,) = possibilities
+            else:
+                # Inferring something else produces too many issues for now.
+                solution = AnyValue(AnySource.inference)
+            return {self.typevar: solution}
+        return {self.typevar: left}
+
     def get_fallback_value(self) -> Value:
-        # TODO: support bounds and bases here to do something smarter
+        if self.bound is not None:
+            return self.bound
+        elif self.constraints:
+            return unite_values(*self.constraints)
         return AnyValue(AnySource.inference)
 
+    def get_type_value(self) -> Value:
+        return self.get_fallback_value().get_type_value()
+
+    def __str__(self) -> str:
+        if self.bound is not None:
+            return f"{self.typevar} <: {self.bound}"
+        elif self.constraints:
+            constraints = ", ".join(map(str, self.constraints))
+            return f"{self.typevar} in ({constraints})"
+        return str(self.typevar)
+
 
 class Extension:
     """An extra piece of information about a type that can be stored in
diff --git a/tox.ini b/tox.ini
index 5dd79de..80d75f1 100644
--- a/tox.ini
+++ b/tox.ini
@@ -12,7 +12,7 @@ commands =
 
 [testenv:black]
 deps =
-    black==21.4b1
+    black==21.12b0
 
 commands =
     black --check pyanalyze/


Test patch:
diff --git a/pyanalyze/test_annotations.py b/pyanalyze/test_annotations.py
index 966a0e5..7c7c928 100644
--- a/pyanalyze/test_annotations.py
+++ b/pyanalyze/test_annotations.py
@@ -12,6 +12,7 @@ from .value import (
     MultiValuedValue,
     NewTypeValue,
     SequenceIncompleteValue,
+    TypeVarValue,
     TypedDictValue,
     TypedValue,
     SubclassValue,
@@ -839,6 +840,104 @@ class TestCallable(TestNameCheckVisitorBase):
             )
 
 
+class TestTypeVar(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_bound(self):
+        from typing import TypeVar
+
+        IntT = TypeVar("IntT", bound=int)
+
+        def f(x: IntT) -> IntT:
+            assert_is_value(x, TypeVarValue(IntT, bound=TypedValue(int)))
+            print(x + 1)
+            return x
+
+        def capybara():
+            assert_is_value(f(1), KnownValue(1))
+            assert_is_value(f(True), KnownValue(True))
+            x = f("")  # E: incompatible_argument
+            assert_is_value(x, AnyValue(AnySource.error))
+
+    @assert_passes()
+    def test_constraint(self):
+        from typing import TypeVar, Union
+
+        AnyStr = TypeVar("AnyStr", bytes, str)
+
+        def whatever(x: Union[str, bytes]):
+            pass
+
+        def f(x: AnyStr) -> AnyStr:
+            print(x.title())
+            whatever(x)
+            return x
+
+        def capybara(s: str, b: bytes, sb: Union[str, bytes], unannotated):
+            assert_is_value(f("x"), TypedValue(str))
+            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
+            f(3)  # E: incompatible_argument
+            assert_is_value(f(unannotated), AnyValue(AnySource.inference))
+
+    @assert_passes()
+    def test_constraint_in_typeshed(self):
+        import re
+
+        def capybara():
+            assert_is_value(re.escape("x"), TypedValue(str))
+
+    @assert_passes()
+    def test_callable_compatibility(self):
+        from typing import TypeVar, Callable, Union
+
+        AnyStr = TypeVar("AnyStr", bytes, str)
+        IntT = TypeVar("IntT", bound=int)
+
+        def want_anystr_func(
+            f: Callable[[AnyStr], AnyStr], s: Union[str, bytes]
+        ) -> str:
+            if isinstance(s, str):
+                assert_is_value(f(s), TypedValue(str))
+            else:
+                assert_is_value(f(s), TypedValue(bytes))
+            return ""
+
+        def want_bounded_func(f: Callable[[IntT], IntT], i: int) -> None:
+            assert_is_value(f(True), KnownValue(True))
+            assert_is_value(f(i), TypedValue(int))
+
+        def want_str_func(f: Callable[[str], str]):
+            assert_is_value(f("x"), TypedValue(str))
+
+        def anystr_func(s: AnyStr) -> AnyStr:
+            return s
+
+        def int_func(i: IntT) -> IntT:
+            return i
+
+        def capybara():
+            want_anystr_func(anystr_func, "x")
+            want_anystr_func(int_func, "x")  # E: incompatible_argument
+            want_bounded_func(int_func, 0)
+            want_bounded_func(anystr_func, 1)  # E: incompatible_argument
+            want_str_func(anystr_func)
+            want_str_func(int_func)  # E: incompatible_argument
+
+    @assert_passes()
+    def test_getitem(self):
+        from typing import Any, Dict, TypeVar, Iterable
+
+        T = TypeVar("T", bound=Dict[str, Any])
+
+        def _fetch_credentials(api: T, credential_names: Iterable[str]) -> T:
+            api_with_credentials = api.copy()
+            for name in credential_names:
+                api_with_credentials[name] = str(api[name])
+            return api_with_credentials
+
+
 class TestParameterTypeGuard(TestNameCheckVisitorBase):
     @assert_passes()
     def test_basic(self):
diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
index a725d91..cd292e7 100644
--- a/pyanalyze/test_name_check_visitor.py
+++ b/pyanalyze/test_name_check_visitor.py
@@ -31,6 +31,7 @@ from .value import (
     ReferencingValue,
     SequenceIncompleteValue,
     TypedValue,
+    TypeVarValue,
     UnboundMethodValue,
     NO_RETURN_VALUE,
     UNINITIALIZED_VALUE,
@@ -134,6 +135,7 @@ def _make_module(code_str: str) -> types.ModuleType:
         SubclassValue=SubclassValue,
         NewTypeValue=NewTypeValue,
         TypedDictValue=TypedDictValue,
+        TypeVarValue=TypeVarValue,
         dump_value=dump_value,
         make_weak=make_weak,
         UNINITIALIZED_VALUE=UNINITIALIZED_VALUE,