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,32 +574,25 @@ 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])
+                    contexts |= ContextSet([generic])
             new_generics.append(contexts)
 
         if not changed:
             # 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>' % (