More preparations for annotated classes

2018-08-27 20:13:35 +02:00
parent 4a7bded98d
commit bd5af5f148
3 changed files with 139 additions and 38 deletions
@@ -221,11 +221,19 @@ class ClassContext(use_metaclass(CachedMetaClass, TreeContext)):
        return ContextName(self, self.tree_node.name)

    def py__getitem__(self, index_context_set, contextualized_node):
-        from jedi.evaluate.context.typing import TypingClassMixin
+        from jedi.evaluate.context.typing import TypingClassMixin, AnnotatedClass
        for cls in self.py__mro__():
            if isinstance(cls, TypingClassMixin):
                #print('ha', self,  list(self.py__mro__()))
                # TODO get the right classes.
-                return ContextSet(self)
+                return ContextSet.from_iterable(
+                    AnnotatedClass(
+                        self.evaluator,
+                        self.parent_context,
+                        self.tree_node,
+                        index_context
+                    )
+                    for index_context in index_context_set
+                )

        return super(ClassContext, self).py__getitem__(index_context_set, contextualized_node)