add CollectArgTypes to get set of argument types

ilevkivskyi · Jul 19, 2022 · Jul 17, 2022 · Jul 17, 2022 · Jul 17, 2022 · Jul 17, 2022
commit 714445ecd33453de2fcc9c2bca8f7ac91a716541
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -40,6 +40,7 @@
     get_proper_types, is_literal_type, TypeAliasType, TypeGuardedType, ParamSpecType,
     OVERLOAD_NAMES, UnboundType
 )
+from mypy.typetraverser import TypeTraverserVisitor
 from mypy.sametypes import is_same_type
 from mypy.messages import (
     MessageBuilder, make_inferred_type_note, append_invariance_notes, pretty_seq,
@@ -919,7 +920,11 @@ def check_func_def(self, defn: FuncItem, typ: CallableType, name: Optional[str])
                         self.fail(message_registry.RETURN_TYPE_CANNOT_BE_CONTRAVARIANT,
                                   typ.ret_type)
 
-                    if not any([isinstance(argtype, TypeVarType) for argtype in typ.arg_types]):
+                    arg_type_visitor = CollectArgTypes()
+                    for argtype in typ.arg_types:
+                        argtype.accept(arg_type_visitor)
+
+                    if typ.ret_type not in arg_type_visitor.arg_types:
                         self.fail(message_registry.RETURN_TYPE_CANNOT_BE_CONTRAVARIANT,
                                   typ.ret_type)
 
@@ -5866,6 +5871,15 @@ class Foo(Enum):
                 and member_type.fallback.type == parent_type.type_object())
 
 
+class CollectArgTypes(TypeTraverserVisitor):
+    """Collects the non-nested argument types in a set."""
+    def __init__(self) -> None:
+        self.arg_types: Set[Instance] = set()
+
+    def visit_type_var(self, t: TypeVarType) -> None:
+        self.arg_types.add(t)
+
+
 @overload
 def conditional_types(current_type: Type,
                       proposed_type_ranges: Optional[List[TypeRange]],

diff --git a/test-data/unit/check-typevar-tuple.test b/test-data/unit/check-typevar-tuple.test
@@ -115,3 +115,23 @@ def f() -> Iterable[T]: # no error
     ...
 
 f()
+
+[case testBoundTypeVar]
+from typing import TypeVar
+
+T = TypeVar('T')
+
+def f(a: T, b: T) -> T: # raise error
+    ...
+
+
+[case testNestedBoundTypeVar]
+from typing import Callable, Union, TypeVar
+
+T = TypeVar('T')
+
+def f(a: Union[int, T], b: str) -> T: # raise error
+    ...
+
+def g(a: Callable[..., T], b: str) -> T: # raise error
+    ...