Get a first typeshed example fully working as intended

jedi/evaluate/context/klass.py

@@ -64,9 +64,7 @@ class ClassName(TreeNameDefinition):
 
     @iterator_to_context_set
     def infer(self):
-        # TODO this _name_to_types might get refactored and be a part of the
-        # parent class. Once it is, we can probably just overwrite method to
-        # achieve this.
+        # We're using a different context to infer, so we cannot call super().
         from jedi.evaluate.syntax_tree import tree_name_to_contexts
         inferred = tree_name_to_contexts(
             self.parent_context.evaluator, self._name_context, self.tree_name)

jedi/plugins/typeshed.py

@@ -5,7 +5,8 @@ from pkg_resources import resource_filename
 from jedi._compatibility import FileNotFoundError
 from jedi.plugins.base import BasePlugin
 from jedi.evaluate.cache import evaluator_function_cache
-from jedi.evaluate.base_context import Context, ContextSet, NO_CONTEXTS
+from jedi.cache import memoize_method
+from jedi.evaluate.base_context import ContextSet, iterator_to_context_set
 from jedi.evaluate.filters import AbstractTreeName, ParserTreeFilter, \
     TreeNameDefinition
 from jedi.evaluate.context import ModuleContext, FunctionContext, ClassContext
@@ -70,6 +71,29 @@ def _load_stub(evaluator, path):
     return evaluator.parse(path=path, cache=True)
 
 
+def _merge_modules(context_set, stub_context):
+    if not context_set:
+        # If there are no results for normal modules, just
+        # use a normal context for stub modules and don't
+        # merge the actual module contexts with stubs.
+        yield stub_context
+        return
+
+    for context in context_set:
+        # TODO what about compiled?
+        if isinstance(context, ModuleContext):
+            yield ModuleStubContext(
+                context.evaluator,
+                stub_context,
+                context.tree_node,
+                context._path,
+                context.code_lines
+            )
+        else:
+            # TODO do we want this?
+            yield context
+
+
 class TypeshedPlugin(BasePlugin):
     _version_cache = {}
 
@@ -120,24 +144,12 @@ class TypeshedPlugin(BasePlugin):
                         # TODO maybe empty cache?
                         pass
                     else:
-                        code_lines = []
-                        args = (
-                            evaluator,
-                            stub_module_node,
-                            path,
-                            code_lines,
+                        # TODO use code_lines
+                        stub_module_context = ModuleContext(
+                            evaluator, stub_module_node, path, code_lines=[]
                         )
-                        if not context_set:
-                            # If there are no results for normal modules, just
-                            # use a normal context for stub modules and don't
-                            # merge the actual module contexts with stubs.
-                            return ModuleContext(*args)
                         return ContextSet.from_iterable(
-                            ModuleStubContext(
-                                *args,
-                                context,
-                                parent_module_context,
-                            ) for context in context_set
+                            _merge_modules(context_set, stub_module_context)
                         )
             # If no stub is found, just return the default.
             return context_set
@@ -152,34 +164,47 @@ class StubName(TreeNameDefinition):
     """
 
     def __init__(self, parent_context, tree_name, stub_parent_context, stub_tree_name):
-        super(StubName, self).__init__(parent_context.actual_context, tree_name)
+        super(StubName, self).__init__(parent_context, tree_name)
         self._stub_parent_context = stub_parent_context
         self._stub_tree_name = stub_tree_name
 
+    @memoize_method
+    @iterator_to_context_set
     def infer(self):
-        def iterate(contexts):
-            for c in contexts:
-                if isinstance(c, FunctionContext):
-                    yield FunctionStubContext(
-                        c.evaluator,
-                        c.parent_context,
-                        c.tree_node,
-                    )
-                elif isinstance(c, ClassContext):
-                    yield ClassStubContext(
-                        c.evaluator,
-                        c.parent_context,
-                        c.tree_node
-                    )
-                else:
-                    yield c
-
-        contexts = tree_name_to_contexts(
+        actual_contexts = super(StubName, self).infer()
+        stub_contexts = tree_name_to_contexts(
             self.parent_context.evaluator,
             self._stub_parent_context,
             self._stub_tree_name
         )
-        return ContextSet.from_iterable(iterate(contexts))
+
+        if not actual_contexts:
+            for c in stub_contexts:
+                yield c
+
+        for actual_context in actual_contexts:
+            for stub_context in stub_contexts:
+                if isinstance(stub_context, FunctionContext) \
+                        and isinstance(actual_context, FunctionContext):
+                    yield FunctionStubContext(
+                        actual_context.evaluator,
+                        stub_context,
+                        actual_context.parent_context,
+                        actual_context.tree_node,
+                    )
+                elif isinstance(stub_context, ClassContext) \
+                        and isinstance(actual_context, ClassContext):
+                    yield ClassStubContext(
+                        actual_context.evaluator,
+                        stub_context,
+                        actual_context.parent_context,
+                        actual_context.tree_node,
+                    )
+                else:
+                    yield stub_context
+
+            if not stub_contexts:
+                yield actual_context
 
 
 class StubParserTreeFilter(ParserTreeFilter):
@@ -201,16 +226,17 @@ class StubParserTreeFilter(ParserTreeFilter):
             # 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 StubName).
-            if not found_actual_names:
+            if not len(found_actual_names):
                 yield TreeNameDefinition(self.context, name)
-            for non_stub_name in found_actual_names:
-                assert isinstance(non_stub_name, AbstractTreeName), non_stub_name
-                yield self.name_class(
-                    non_stub_name.parent_context,
-                    non_stub_name.tree_name,
-                    self.context,
-                    name,
-                )
+            else:
+                for non_stub_name in found_actual_names:
+                    assert isinstance(non_stub_name, AbstractTreeName), non_stub_name
+                    yield self.name_class(
+                        non_stub_name.parent_context,
+                        non_stub_name.tree_name,
+                        self.context,
+                        name,
+                    )
 
     def _is_name_reachable(self, name):
         if not super(StubParserTreeFilter, self)._is_name_reachable(name):
@@ -226,55 +252,18 @@ class StubParserTreeFilter(ParserTreeFilter):
         return True
 
 
-class StubProxy(object):
-    def __init__(self, stub_context, parent_context):
-        self._stub_context = stub_context
-        self._parent_context = parent_context
-
-    def get_filters(self, *args, **kwargs):
-        for f in self._stub_context.get_filters(*args, **kwargs):
-            yield StubFilterWrapper(f)
-
-    # We have to overwrite everything that has to do with trailers, name
-    # lookups and filters to make it possible to route name lookups towards
-    # compiled objects and the rest towards tree node contexts.
-    def py__getattribute__(self, *args, **kwargs):
-        return self._stub_context.py__getattribute__(*args, **kwargs)
-        #context_results = self._context.py__getattribute__(
-        #    *args, **kwargs
-        #)
-        typeshed_results = list(self._stub_context.py__getattribute__(
-            *args, **kwargs
-        ))
-        if not typeshed_results:
-            return NO_CONTEXTS
-
-        return ContextSet.from_iterable(
-            StubProxy(c) for c in typeshed_results
-        )
-
-    def get_root_context(self):
-        if self._parent_context is None:
-            return self
-
-        return self._parent_context.get_root_context()
-
-    def __getattr__(self, name):
-        return getattr(self._stub_context, name)
-
-    def __repr__(self):
-        return '<%s: %s>' % (type(self).__name__, self._stub_context)
+class _MixedStubContextMixin(object):
+    """
+    Mixes the actual contexts with the stub module contexts.
+    """
+    def __init__(self, evaluator, stub_context, *args, **kwargs):
+        super(_MixedStubContextMixin, self).__init__(evaluator, *args, **kwargs)
+        self.stub_context = stub_context
 
 
-class ModuleStubContext(ModuleContext):
-    def __init__(self, evaluator, stub_module_node, path, code_lines,
-                 actual_context, parent_module_context):
-        super(ModuleStubContext, self).__init__(evaluator, stub_module_node, path, code_lines),
-        self._parent_module_context = parent_module_context
-        self.actual_context = actual_context
-
+class _StubContextFilterMixin(_MixedStubContextMixin):
     def get_filters(self, search_global, until_position=None, origin_scope=None):
-        filters = super(ModuleStubContext, self).get_filters(
+        filters = super(_StubContextFilterMixin, self).get_filters(
             search_global, until_position, origin_scope
         )
         yield StubParserTreeFilter(
@@ -282,7 +271,7 @@ class ModuleStubContext(ModuleContext):
             # and wrap it.
             next(filters),
             self.evaluator,
-            context=self,
+            context=self.stub_context,
             until_position=until_position,
             origin_scope=origin_scope,
             search_global=search_global,
@@ -291,9 +280,13 @@ class ModuleStubContext(ModuleContext):
             yield f
 
 
-class ClassStubContext(ClassContext):
+class ModuleStubContext(_StubContextFilterMixin, ModuleContext):
     pass
 
 
-class FunctionStubContext(FunctionContext):
+class ClassStubContext(_StubContextFilterMixin, ClassContext):
+    pass
+
+
+class FunctionStubContext(_MixedStubContextMixin, FunctionContext):
     pass

test/test_plugin/test_stub.py

@@ -48,16 +48,23 @@ def test_function(Script):
 
 
 def test_class(Script):
-    def_, = Script('import threading; threading.Lock').goto_definitions()
+    def_, = Script('import threading; threading.Thread').goto_definitions()
     context = def_._name._context
     assert isinstance(context, typeshed.ClassStubContext), context
 
 
 def test_instance(Script):
-    s = Script('import threading; threading.Lock()')
+    s = Script('import threading; threading.Thread()')
+
+
+def test_class_function(Script):
+    def_, = Script('import threading; threading.Thread.getName').goto_definitions()
+    context = def_._name._context
+    assert isinstance(context, typeshed.FunctionStubContext), context
 
 
 def test_method(Script):
-    s = Script('import threading; threading.Lock().locked')
-
-
+    code = 'import threading; threading.Thread().getName'
+    def_, = Script(code).goto_definitions()
+    context = def_._name._context
+    assert isinstance(context, typeshed.ClassStubContext), context