Use TreeContext in a good way

jedi/evaluate/base_context.py

@@ -199,9 +199,10 @@ def iterate_contexts(contexts, contextualized_node=None, is_async=False):
 
 
 class TreeContext(Context):
-    def __init__(self, evaluator, parent_context=None):
+    def __init__(self, evaluator, parent_context, tree_node):
         super(TreeContext, self).__init__(evaluator, parent_context)
         self.predefined_names = {}
+        self.tree_node = tree_node
 
     def __repr__(self):
         return '<%s: %s>' % (self.__class__.__name__, self.tree_node)

jedi/evaluate/compiled/context.py

@@ -55,7 +55,7 @@ class CompiledObject(Context):
             return FunctionContext(
                 self.evaluator,
                 parent_context=self.parent_context,
-                funcdef=self.tree_node
+                tree_node=self.tree_node
             ).py__call__(params)
         if self.access_handle.is_class():
             from jedi.evaluate.context import CompiledInstance

jedi/evaluate/context/function.py

@@ -44,11 +44,6 @@ class FunctionContext(use_metaclass(CachedMetaClass, TreeContext)):
     """
     api_type = u'function'
 
-    def __init__(self, evaluator, parent_context, funcdef):
-        """ This should not be called directly """
-        super(FunctionContext, self).__init__(evaluator, parent_context)
-        self.tree_node = funcdef
-
     def get_filters(self, search_global, until_position=None, origin_scope=None):
         if search_global:
             yield ParserTreeFilter(
@@ -127,9 +122,12 @@ class FunctionExecutionContext(TreeContext):
     function_execution_filter = FunctionExecutionFilter
 
     def __init__(self, evaluator, parent_context, function_context, var_args):
-        super(FunctionExecutionContext, self).__init__(evaluator, parent_context)
+        super(FunctionExecutionContext, self).__init__(
+            evaluator,
+            parent_context,
+            function_context.tree_node,
+        )
         self.function_context = function_context
-        self.tree_node = function_context.tree_node
         self.var_args = var_args
 
     @evaluator_method_cache(default=NO_CONTEXTS)

jedi/evaluate/context/instance.py

@@ -28,6 +28,7 @@ class InstanceFunctionExecution(BaseInstanceFunctionExecution):
 
         super(InstanceFunctionExecution, self).__init__(
             instance, parent_context, function_context, var_args)
+        print(parent_context, self.parent_context, function_context)
 
 
 class AnonymousInstanceFunctionExecution(BaseInstanceFunctionExecution):

jedi/evaluate/context/iterable.py

@@ -83,10 +83,6 @@ class CompForContext(TreeContext):
     def from_comp_for(cls, parent_context, comp_for):
         return cls(parent_context.evaluator, parent_context, comp_for)
 
-    def __init__(self, evaluator, parent_context, comp_for):
-        super(CompForContext, self).__init__(evaluator, parent_context)
-        self.tree_node = comp_for
-
     def get_node(self):
         return self.tree_node
 

jedi/evaluate/context/klass.py

@@ -91,10 +91,6 @@ class ClassContext(use_metaclass(CachedMetaClass, TreeContext)):
     """
     api_type = u'class'
 
-    def __init__(self, evaluator, parent_context, classdef):
-        super(ClassContext, self).__init__(evaluator, parent_context=parent_context)
-        self.tree_node = classdef
-
     @evaluator_method_cache(default=())
     def py__mro__(self):
         def add(cls):

jedi/evaluate/context/module.py

@@ -43,8 +43,11 @@ class ModuleContext(TreeContext):
     parent_context = None
 
     def __init__(self, evaluator, module_node, path, code_lines):
-        super(ModuleContext, self).__init__(evaluator, parent_context=None)
-        self.tree_node = module_node
+        super(ModuleContext, self).__init__(
+            evaluator,
+            parent_context=None,
+            tree_node=module_node
+        )
         self._path = path
         self.code_lines = code_lines
 

jedi/evaluate/context/namespace.py

@@ -4,7 +4,7 @@ from itertools import chain
 from jedi.evaluate.cache import evaluator_method_cache
 from jedi.evaluate import imports
 from jedi.evaluate.filters import DictFilter, AbstractNameDefinition, ContextNameMixin
-from jedi.evaluate.base_context import TreeContext, ContextSet
+from jedi.evaluate.base_context import Context
 
 
 class ImplicitNSName(ContextNameMixin, AbstractNameDefinition):
@@ -17,7 +17,7 @@ class ImplicitNSName(ContextNameMixin, AbstractNameDefinition):
         self.string_name = string_name
 
 
-class ImplicitNamespaceContext(TreeContext):
+class ImplicitNamespaceContext(Context):
     """
     Provides support for implicit namespace packages
     """

jedi/evaluate/syntax_tree.py

@@ -580,13 +580,13 @@ def _apply_decorators(context, node):
         decoratee_context = ClassContext(
             context.evaluator,
             parent_context=context,
-            classdef=node
+            tree_node=node
         )
     else:
         decoratee_context = FunctionContext(
             context.evaluator,
             parent_context=context,
-            funcdef=node
+            tree_node=node
         )
     initial = values = ContextSet(decoratee_context)
     for dec in reversed(node.get_decorators()):