Better wrapping of BoundMethod

jedi/evaluate/context/function.py

@@ -37,7 +37,7 @@ class LambdaName(AbstractNameDefinition):
         return ContextSet([self._lambda_context])
 
 
-class AbstractFunction(TreeContext):
+class FunctionMixin(object):
     api_type = u'function'
 
     def get_filters(self, search_global=False, until_position=None, origin_scope=None):
@@ -65,13 +65,16 @@ class AbstractFunction(TreeContext):
         return ContextName(self, self.tree_node.name)
 
     def get_function_execution(self, arguments=None):
-        raise NotImplementedError
+        if arguments is None:
+            arguments = AnonymousArguments()
+
+        return FunctionExecutionContext(self.evaluator, self.parent_context, self, arguments)
 
     def py__name__(self):
         return self.name.string_name
 
 
-class FunctionContext(use_metaclass(CachedMetaClass, AbstractFunction)):
+class FunctionContext(use_metaclass(CachedMetaClass, FunctionMixin, TreeContext)):
     """
     Needed because of decorators. Decorators are evaluated here.
     """
@@ -116,12 +119,6 @@ class FunctionContext(use_metaclass(CachedMetaClass, AbstractFunction)):
             return NO_CONTEXTS
         return function_execution.infer()
 
-    def get_function_execution(self, arguments=None):
-        if arguments is None:
-            arguments = AnonymousArguments()
-
-        return FunctionExecutionContext(self.evaluator, self.parent_context, self, arguments)
-
     def py__class__(self):
         return compiled.get_special_object(self.evaluator, u'FUNCTION_CLASS')
 
@@ -135,10 +132,10 @@ class FunctionContext(use_metaclass(CachedMetaClass, AbstractFunction)):
 class MethodContext(FunctionContext):
     def __init__(self, evaluator, class_context, *args, **kwargs):
         super(MethodContext, self).__init__(evaluator, *args, **kwargs)
-        self.class_context = class_context
+        self._class_context = class_context
 
     def get_default_param_context(self):
-        return self.class_context
+        return self._class_context
 
 
 class FunctionExecutionContext(TreeContext):

jedi/evaluate/context/instance.py

@@ -5,13 +5,13 @@ from jedi import settings
 from jedi.evaluate import compiled
 from jedi.evaluate import filters
 from jedi.evaluate.base_context import Context, NO_CONTEXTS, ContextSet, \
-    iterator_to_context_set
+    iterator_to_context_set, ContextWrapper
 from jedi.evaluate.lazy_context import LazyKnownContext, LazyKnownContexts
 from jedi.evaluate.cache import evaluator_method_cache
 from jedi.evaluate.arguments import AbstractArguments, AnonymousArguments, \
     ValuesArguments, TreeArgumentsWrapper
 from jedi.evaluate.context.function import FunctionExecutionContext, \
-    FunctionContext, AbstractFunction, OverloadedFunctionContext
+    FunctionContext, FunctionMixin, OverloadedFunctionContext
 from jedi.evaluate.context.klass import ClassContext, apply_py__get__, \
     py__mro__, ClassFilter
 from jedi.evaluate.context import iterable
@@ -355,16 +355,11 @@ class CompiledInstanceClassFilter(filters.AbstractFilter):
         ]
 
 
-class BoundMethod(AbstractFunction):
+class BoundMethod(FunctionMixin, ContextWrapper):
     def __init__(self, instance, klass, function):
-        super(BoundMethod, self).__init__(
-            function.evaluator,
-            function.parent_context,
-            function.tree_node,
-        )
+        super(BoundMethod, self).__init__(function)
         self.instance = instance
         self.class_context = klass
-        self._function = function
 
     def py__class__(self):
         return compiled.get_special_object(self.evaluator, u'BOUND_METHOD_CLASS')
@@ -378,33 +373,30 @@ class BoundMethod(AbstractFunction):
     def get_function_execution(self, arguments=None):
         arguments = self._get_arguments(arguments)
 
-        if isinstance(self._function, compiled.CompiledObject):
+        if isinstance(self._wrapped_context, compiled.CompiledObject):
             # This is kind of weird, because it's coming from a compiled object
             # and we're not sure if we want that in the future.
             return FunctionExecutionContext(
                 self.evaluator, self.parent_context, self, arguments
             )
 
-        return self._function.get_function_execution(arguments)
-
-    def get_default_param_context(self):
-        return self.class_context
+        return super(BoundMethod, self).get_function_execution(arguments)
 
     def py__call__(self, arguments):
-        if isinstance(self._function, OverloadedFunctionContext):
-            return self._function.py__call__(self._get_arguments(arguments))
+        if isinstance(self._wrapped_context, OverloadedFunctionContext):
+            return self._wrapped_context.py__call__(self._get_arguments(arguments))
         function_execution = self.get_function_execution(arguments)
         return function_execution.infer()
 
     def get_matching_functions(self, arguments):
-        for func in self._function.get_matching_functions(arguments):
+        for func in self._wrapped_context.get_matching_functions(arguments):
             if func is self:
                 yield self
             else:
                 yield BoundMethod(self.instance, self.class_context, func)
 
     def __repr__(self):
-        return '<%s: %s>' % (self.__class__.__name__, self._function)
+        return '<%s: %s>' % (self.__class__.__name__, self._wrapped_context)
 
 
 class CompiledBoundMethod(compiled.CompiledObject):

jedi/evaluate/finder.py

@@ -185,8 +185,15 @@ class NameFinder(object):
         if not contexts and isinstance(self._name, tree.Name) and \
                 not self._name_context.is_instance():
             flow_scope = self._name
-            base_node = self._name_context.tree_node
-            if base_node.type == 'comp_for':
+            base_nodes = [self._name_context.tree_node]
+            try:
+                stub_node = self._name_context.stub_context.tree_node
+            except AttributeError:
+                pass
+            else:
+                base_nodes.append(stub_node)
+
+            if any(b.type == 'comp_for' for b in base_nodes):
                 return contexts
             while True:
                 flow_scope = get_parent_scope(flow_scope, include_flows=True)
@@ -194,7 +201,7 @@ class NameFinder(object):
                                             self._name, self._position)
                 if n is not None:
                     return n
-                if flow_scope == base_node:
+                if flow_scope in base_nodes:
                     break
         return contexts