Some more work on the filter merging

jedi/api/classes.py

@@ -14,7 +14,7 @@ from jedi.evaluate import imports
 from jedi.evaluate import compiled
 from jedi.evaluate.imports import ImportName
 from jedi.evaluate.context import instance
-from jedi.evaluate.context import ClassContext, FunctionExecutionContext
+from jedi.evaluate.context import FunctionExecutionContext
 from jedi.plugins.typeshed import StubOnlyModuleContext
 from jedi.api.keywords import KeywordName
 
@@ -339,7 +339,7 @@ class BaseDefinition(object):
                 # there's no better solution.
                 inferred = names[0].infer()
                 param_names = get_param_names(next(iter(inferred)))
-                if isinstance(context, ClassContext):
+                if context.is_class():
                     param_names = param_names[1:]
                 return param_names
             elif isinstance(context, compiled.CompiledObject):

jedi/evaluate/base_context.py

@@ -18,8 +18,6 @@ from jedi.evaluate.cache import evaluator_as_method_param_cache
 
 
 class HelperContextMixin:
-    tree_node = None
-
     @classmethod
     @evaluator_as_method_param_cache()
     def create_cached(cls, *args, **kwargs):
@@ -62,6 +60,8 @@ class HelperContextMixin:
         return False
 
     def is_same_class(self, class2):
+        if isinstance(class2, ContextWrapper):
+            class2 = class2._wrapped_context
         # Class matching should prefer comparisons that are not this function.
         if type(class2).is_same_class != HelperContextMixin.is_same_class:
             return class2.is_same_class(self)
@@ -76,6 +76,7 @@ class Context(HelperContextMixin, BaseContext):
     """
     To be defined by subclasses.
     """
+    tree_node = None
 
     @property
     def api_type(self):

jedi/evaluate/compiled/context.py

@@ -373,6 +373,9 @@ class CompiledObjectFilter(AbstractFilter):
     def _create_name(self, name):
         return self.name_class(self._evaluator, self._compiled_object, name)
 
+    def __repr__(self):
+        return "<%s: %s>" % (self.__class__.__name__, self._compiled_object)
+
 
 docstr_defaults = {
     'floating point number': u'float',

jedi/evaluate/context/klass.py

@@ -63,6 +63,7 @@ def apply_py__get__(context, base_context):
 @evaluator_method_cache(default=())
 def py__mro__(context):
     try:
+        # TODO is this really needed?
         method = context.py__mro__
     except AttributeError:
         pass

jedi/evaluate/pep0484.py

@@ -26,7 +26,6 @@ from parso import ParserSyntaxError, parse
 from jedi._compatibility import force_unicode
 from jedi.evaluate.cache import evaluator_method_cache
 from jedi.evaluate.base_context import ContextSet, NO_CONTEXTS
-from jedi.evaluate.context import ClassContext
 from jedi.evaluate.context.typing import TypeVar, AnnotatedClass, \
     AbstractAnnotatedClass
 from jedi.evaluate.helpers import is_string
@@ -232,12 +231,15 @@ def infer_return_types(function_execution_context):
         However, the iterator is defined as Iterator[_T_co], which means it has
         a different type var name.
         """
-        if isinstance(context, ClassContext):
+        try:
+            func = context.list_type_vars
+        except AttributeError:
+            return type_var_dict
+        else:
             return {
                 to.py__name__(): type_var_dict.get(from_.py__name__(), NO_CONTEXTS)
-                for from_, to in zip(unknown_type_vars, context.list_type_vars())
+                for from_, to in zip(unknown_type_vars, func())
             }
-        return type_var_dict
 
     return ContextSet(
         ann.define_generics(type_var_dict)

jedi/plugins/typeshed.py

@@ -9,7 +9,7 @@ from jedi.parser_utils import get_call_signature_for_any
 from jedi.evaluate.base_context import ContextSet, iterator_to_context_set, \
     ContextWrapper
 from jedi.evaluate.filters import AbstractTreeName, ParserTreeFilter, \
-    NameWrapper, AbstractFilter
+    NameWrapper, AbstractFilter, TreeNameDefinition
 from jedi.evaluate.context import ModuleContext, FunctionContext, \
     ClassContext
 from jedi.evaluate.context.typing import TypingModuleFilterWrapper, \
@@ -196,7 +196,7 @@ class NameWithStubMixin(object):
                         actual_context.parent_context,
                         actual_context.tree_node,
                     )
-                elif isinstance(stub_context, ClassContext) \
+                elif isinstance(stub_context, StubOnlyClass) \
                         and isinstance(actual_context, ClassContext):
                     yield StubClassContext(
                         actual_context.evaluator,
@@ -211,9 +211,18 @@ class NameWithStubMixin(object):
                 yield actual_context
 
 
-class NameWithStub(NameWithStubMixin, NameWrapper):
+class StubOnlyName(TreeNameDefinition):
+    def infer(self):
+        inferred = super(StubOnlyName, self).infer()
+        return [
+            StubOnlyClass(c) if isinstance(c, ClassContext) else c
+            for c in inferred
+        ]
+
+
+class StubName(NameWithStubMixin, NameWrapper):
     def __init__(self, non_stub_name, stub_name):
-        super(NameWithStub, self).__init__(non_stub_name)
+        super(StubName, self).__init__(non_stub_name)
         self._stub_name = stub_name
 
     def _get_actual_contexts(self):
@@ -234,6 +243,8 @@ class CompiledNameWithStub(NameWithStubMixin, NameWrapper):
 
 
 class StubOnlyFilter(ParserTreeFilter):
+    name_class = StubOnlyName
+
     def __init__(self, *args, **kwargs):
         self._search_global = kwargs.pop('search_global')  # Python 2 :/
         super(StubOnlyFilter, self).__init__(*args, **kwargs)
@@ -261,23 +272,11 @@ class StubFilter(AbstractFilter):
         self._stub_filters = stub_filters
 
     def get(self, name):
-        non_stub_names = self._get_names_from_filter(self._non_stub_filters)
+        non_stub_names = self._get_names_from_filters(self._non_stub_filters, name)
-        stub_names = self._get_names_from_filter(self._stub_filters)
+        stub_names = self._get_names_from_filters(self._stub_filters, name)
-        if stub_names and non_stub_names:
-            # This is the case where we need to merge, because the names are
-            # contained in both filters.
         return self._merge_names(non_stub_names, stub_names)
-        return stub_names or non_stub_names
 
     def values(self):
-        used_stub_names = set()
-        result_names = []
-        for key_name, names in self._used_names.items():
-            found_names = self._convert_names(self._filter(names))
-            if found_names:
-                result_names += found_names
-                used_stub_names.add(key_name)
-
         name_dict = {}
         for non_stub_filter in self._non_stub_filters:
             for name in non_stub_filter.values():
@@ -286,7 +285,7 @@ class StubFilter(AbstractFilter):
         # Try to match the names of stubs with non-stubs. If there's no
         # match, just use the stub name. The user will be directed there
         # for all API accesses. Otherwise the user will be directed to the
-        # non-stub positions (see NameWithStub).
+        # non-stub positions (see StubName).
         for stub_filter in self._stub_filters:
             for stub_name in stub_filter.values():
                 merged_names = self._merge_names(
@@ -305,18 +304,25 @@ class StubFilter(AbstractFilter):
 
     @to_list
     def _merge_names(self, names, stub_names):
+        if not stub_names:
+            return names
+        if not names:
+            if isinstance(self._stub_filters[0].context, TypingModuleWrapper):
+                return [TypingModuleName(n) for n in stub_names]
+            return stub_names
+
+        result = []
+        # The names are contained in both filters.
         for name in names:
-            if isinstance(self._non_stub_filters[0].context, TypingModuleWrapper):
-                name = TypingModuleName(name)
             for stub_name in stub_names:
-                if isinstance(stub_name, CompiledName):
+                if isinstance(self._stub_filters[0].context, TypingModuleWrapper):
-                    yield CompiledNameWithStub(
+                    stub_name = TypingModuleName(stub_name)
-                        name,
+
-                        stub_name,
+                if isinstance(name, CompiledName):
-                    )
+                    result.append(CompiledNameWithStub(name, stub_name))
                 else:
-                    assert isinstance(stub_name, AbstractTreeName), stub_name
+                    result.append(StubName(name, stub_name))
-                    yield NameWithStub(name, stub_name)
+        return result
 
 
 class _MixedStubContextMixin(object):
@@ -334,15 +340,13 @@ class _StubContextFilterMixin(_MixedStubContextMixin):
         filters = super(_StubContextFilterMixin, self).get_filters(
             search_global, until_position, origin_scope, **kwargs
         )
-        yield StubFilter(
+        yield self.stub_context.get_stub_only_filter(
             # Take the first filter, which is here to filter module contents
             # and wrap it.
             [next(filters)],
-            self.evaluator,
+            search_global=search_global,
-            context=self.stub_context,
             until_position=until_position,
             origin_scope=origin_scope,
-            search_global=search_global,
         )
         for f in filters:
             yield f
@@ -355,7 +359,7 @@ class StubModuleContext(_StubContextFilterMixin, ModuleContext):
 class StubClassContext(_StubContextFilterMixin, ClassContext):
     def __getattribute__(self, name):
         if name in ('py__getitem__', 'py__simple_getitem__', 'py__bases__',
-                    'execute_annotation'):
+                    'execute_annotation', 'get_stub_only_filter'):
             # getitem is always done in the stub class.
             return getattr(self.stub_context, name)
         return super(StubClassContext, self).__getattribute__(name)
@@ -367,7 +371,24 @@ class StubFunctionContext(_MixedStubContextMixin, FunctionContext):
         return super().get_function_execution(arguments, tree_node=self.stub_context.tree_node)
 
 
-class StubOnlyModuleContext(ModuleContext):
+class _StubOnlyContext(object):
+    def _get_stub_only_filters(self, **filter_kwargs):
+        return [StubOnlyFilter(
+            self.evaluator,
+            context=self,
+            **filter_kwargs
+        )]
+
+    def get_stub_only_filter(self, non_stub_filters, **filter_kwargs):
+        # Here we remap the names from stubs to the actual module. This is
+        # important if type inferences is needed in that module.
+        return StubFilter(
+            non_stub_filters,
+            self._get_stub_only_filters(**filter_kwargs),
+        )
+
+
+class StubOnlyModuleContext(_StubOnlyContext, ModuleContext):
     def __init__(self, non_stub_context_set, *args, **kwargs):
         super(StubOnlyModuleContext, self).__init__(*args, **kwargs)
         self.non_stub_context_set = non_stub_context_set
@@ -376,6 +397,13 @@ class StubOnlyModuleContext(ModuleContext):
         for context in self.non_stub_context_set:
             yield next(context.get_filters(search_global=False))
 
+    def _get_stub_only_filters(self, search_global, **filter_kwargs):
+        stub_filters = super(StubOnlyModuleContext, self)._get_stub_only_filters(
+            search_global=search_global, **filter_kwargs
+        )
+        stub_filters += self.iter_star_filters(search_global=search_global)
+        return stub_filters
+
     def get_filters(self, search_global=False, until_position=None,
                     origin_scope=None, **kwargs):
         filters = super(StubOnlyModuleContext, self).get_filters(
@@ -383,28 +411,18 @@ class StubOnlyModuleContext(ModuleContext):
         )
         next(filters)  # Ignore the first filter and replace it with our own
 
-        stub_filters = [StubOnlyFilter(
+        yield self.get_stub_only_filter(
-            self.evaluator,
-            context=self,
-            until_position=until_position,
-            origin_scope=origin_scope,
-            search_global=search_global,
-        )]
-        stub_filters += self.
-
-        # Here we remap the names from stubs to the actual module. This is
-        # important if type inferences is needed in that module.
-        yield StubFilter(
             list(self._get_first_non_stub_filters()),
-            self.evaluator,
+            search_global=search_global,
-            context=self,
             until_position=until_position,
             origin_scope=origin_scope,
-            search_global=search_global,
         )
         for f in filters:
             yield f
 
+class StubOnlyClass(_StubOnlyContext, ContextWrapper):
+    pass
+
 
 class _StubContextWithCompiled(ContextWrapper):
     def __init__(self, stub_context, compiled_context):