Fix stub function inferrals

jedi/evaluate/context/instance.py

@@ -283,7 +283,7 @@ class TreeInstance(AbstractInstanceContext):
                 continue
 
             all_annotations = py__annotations__(execution.tree_node)
-            defined = self.class_context.define_generics(
+            defined, = self.class_context.define_generics(
                 infer_type_vars_for_execution(execution, all_annotations),
             )
             debug.dbg('Inferred instance context as %s', defined, color='BLUE')

jedi/evaluate/context/klass.py

@@ -272,11 +272,11 @@ class ClassContext(use_metaclass(CachedMetaClass, ClassMixin, FunctionAndClassBa
                 yield type_var_dict.get(type_var.py__name__(), NO_CONTEXTS)
 
         if type_var_dict:
-            return GenericClass(
+            return ContextSet([GenericClass(
                 self,
                 generics=tuple(remap_type_vars())
-            )
-        return self
+            )])
+        return ContextSet({self})
 
     def get_signatures(self):
         init_funcs = self.py__getattribute__('__init__')

jedi/evaluate/gradual/annotation.py

@@ -203,9 +203,9 @@ def infer_return_types(function_execution_context):
 
     type_var_dict = infer_type_vars_for_execution(function_execution_context, all_annotations)
 
-    return ContextSet(
+    return ContextSet.from_sets(
         ann.define_generics(type_var_dict)
-        if isinstance(ann, AbstractAnnotatedClass) else ann
+        if isinstance(ann, (AbstractAnnotatedClass, TypeVar)) else ann
         for ann in annotation_contexts
     ).execute_annotation()
 

jedi/evaluate/gradual/typing.py

@@ -432,6 +432,16 @@ class TypeVar(_BaseTypingContext):
             lazy.infer() for lazy in self._constraints_lazy_contexts
         )
 
+    def define_generics(self, type_var_dict):
+        try:
+            found = type_var_dict[self.py__name__()]
+        except KeyError:
+            pass
+        else:
+            if found:
+                return found
+        return self._get_classes() or ContextSet({self})
+
     def execute_annotation(self):
         return self._get_classes().execute_annotation()
 
@@ -564,18 +574,11 @@ class AbstractAnnotatedClass(ClassMixin, ContextWrapper):
         for generic_set in self.get_generics():
             contexts = NO_CONTEXTS
             for generic in generic_set:
-                if isinstance(generic, AbstractAnnotatedClass):
-                    new_generic = generic.define_generics(type_var_dict)
-                    contexts |= ContextSet([new_generic])
-                    if new_generic != generic:
+                if isinstance(generic, (AbstractAnnotatedClass, TypeVar)):
+                    result = generic.define_generics(type_var_dict)
+                    contexts |= result
+                    if result != ContextSet({generic}):
                         changed = True
-                else:
-                    if isinstance(generic, TypeVar):
-                        try:
-                            contexts |= type_var_dict[generic.py__name__()]
-                            changed = True
-                        except KeyError:
-                            contexts |= ContextSet([generic])
                 else:
                     contexts |= ContextSet([generic])
             new_generics.append(contexts)
@@ -584,12 +587,12 @@ class AbstractAnnotatedClass(ClassMixin, ContextWrapper):
             # There might not be any type vars that change. In that case just
             # return itself, because it does not make sense to potentially lose
             # cached results.
-            return self
+            return ContextSet([self])
 
-        return GenericClass(
+        return ContextSet([GenericClass(
             self._wrapped_context,
             generics=tuple(new_generics)
-        )
+        )])
 
     def __repr__(self):
         return '<%s: %s%s>' % (

test/completion/pep0484_typing.py

@@ -306,17 +306,21 @@ def union4(x: U[int, str]):
     #? int() str()
     x
 
+# -------------------------
+# Type Vars
+# -------------------------
 
-TYPE_VAR = typing.TypeVar('TYPE_VAR')
+TYPE_VARX = typing.TypeVar('TYPE_VARX')
+TYPE_VAR_CONSTRAINTSX = typing.TypeVar('TYPE_VAR_CONSTRAINTSX', str, int)
 # TODO there should at least be some results.
 #? []
-TYPE_VAR.
-#! ["TYPE_VAR = typing.TypeVar('TYPE_VAR')"]
-TYPE_VAR
+TYPE_VARX.
+#! ["TYPE_VARX = typing.TypeVar('TYPE_VARX')"]
+TYPE_VARX
 
 
-class WithTypeVar(typing.Generic[TYPE_VAR]):
-    def lala(self) -> TYPE_VAR:
+class WithTypeVar(typing.Generic[TYPE_VARX]):
+    def lala(self) -> TYPE_VARX:
         ...
 
 
@@ -324,6 +328,33 @@ def maaan(p: WithTypeVar[int]):
     #? int()
     p.lala()
 
+def in_out1(x: TYPE_VARX) -> TYPE_VARX: ...
+
+#? int()
+in_out1(1)
+#? str()
+in_out1("")
+#? str()
+in_out1(str())
+#?
+in_out1()
+
+def in_out2(x: TYPE_VAR_CONSTRAINTSX) -> TYPE_VAR_CONSTRAINTSX: ...
+
+#? int()
+in_out2(1)
+#? str()
+in_out2("")
+#? str()
+in_out2(str())
+#? str() int()
+in_out2()
+##? str() int()
+in_out2(1.0)
+
+# -------------------------
+# TYPE_CHECKING
+# -------------------------
 
 if typing.TYPE_CHECKING:
     with_type_checking = 1
@@ -352,6 +383,10 @@ def foo(a: typing.List, b: typing.Dict, c: typing.MutableMapping) -> typing.Type
 #? int
 foo()
 
+# -------------------------
+# cast
+# -------------------------
+
 def cast_tests():
     x = 3.0
     y = typing.cast(int, x)