Define generics from a different function

jedi/evaluate/context/instance.py

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

jedi/evaluate/context/iterable.py

@@ -529,6 +529,12 @@ class FakeDict(_DictMixin, _FakeArray):
     def _dict_values(self):
         return ContextSet.from_sets(lazy_context.infer() for lazy_context in self._dct.values())
 
+    def _dict_keys(self):
+        return ContextSet.from_sets(lazy_context.infer() for lazy_context in self.py__iter__())
+
+    def get_mapping_item_contexts(self):
+        return self._dict_keys(), self._dict_values()
+
     def exact_key_items(self):
         return self._dct.items()
 

jedi/evaluate/context/klass.py

@@ -47,7 +47,7 @@ from jedi.evaluate.filters import ParserTreeFilter, TreeNameDefinition, \
     ContextName
 from jedi.evaluate.arguments import unpack_arglist
 from jedi.evaluate.base_context import ContextSet, iterator_to_context_set, \
-    TreeContext
+    TreeContext, NO_CONTEXTS
 
 
 def apply_py__get__(context, base_context):
@@ -274,3 +274,19 @@ class ClassContext(use_metaclass(CachedMetaClass, TreeContext)):
             )
             for index_context in index_context_set
         )
+
+    def define_generics(self, type_var_dict):
+        from jedi.evaluate.context.typing import AnnotatedSubClass
+
+        def remap_type_vars():
+            for type_var in self.list_type_vars():
+                yield type_var_dict.get(type_var.py__name__(), NO_CONTEXTS)
+
+        if type_var_dict:
+            return AnnotatedSubClass(
+                self.evaluator,
+                self.parent_context,
+                self.tree_node,
+                given_types=tuple(remap_type_vars())
+            )
+        return self

jedi/evaluate/filters.py

@@ -7,7 +7,6 @@ from abc import abstractmethod
 from parso.tree import search_ancestor
 
 from jedi._compatibility import use_metaclass, Parameter
-from jedi.cache import memoize_method
 from jedi.evaluate import flow_analysis
 from jedi.evaluate.base_context import ContextSet, Context
 from jedi.parser_utils import get_parent_scope

jedi/evaluate/pep0484.py

@@ -290,21 +290,6 @@ def infer_type_vars_for_execution(execution_context, annotation_dict):
     return annotation_variable_results
 
 
-def define_type_vars(annotation_context, type_var_dict):
-    def remap_type_vars(cls):
-        for type_var in cls.list_type_vars():
-            yield type_var_dict.get(type_var.py__name__(), NO_CONTEXTS)
-
-    if type_var_dict and isinstance(annotation_context, ClassContext):
-        return AnnotatedSubClass(
-            annotation_context.evaluator,
-            annotation_context.parent_context,
-            annotation_context.tree_node,
-            given_types=tuple(remap_type_vars(annotation_context))
-        )
-    return annotation_context
-
-
 def _merge_type_var_dicts(base_dict, new_dict):
     for type_var_name, contexts in new_dict.items():
         try: